diff --git a/migrations/versions/a08ce5e7246a_adds_api_key_field_to_pool_table.py b/migrations/versions/a08ce5e7246a_adds_api_key_field_to_pool_table.py
index 3da561c..b5b04f9 100644
--- a/migrations/versions/a08ce5e7246a_adds_api_key_field_to_pool_table.py
+++ b/migrations/versions/a08ce5e7246a_adds_api_key_field_to_pool_table.py
@@ -6,33 +6,28 @@ Create Date: 2022-12-20 18:10:19.540534
 
 """
 import secrets
+from datetime import datetime
 
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.orm import Session
-
-from app.models.base import Pool
+from sqlalchemy.ext.automap import automap_base
 
 revision = 'a08ce5e7246a'
 down_revision = '6a59928efeb7'
 branch_labels = None
 depends_on = None
 
-
-# class Pool(db.Model):
-#     id = db.Column(db.Integer, primary_key=True)
-#     description = db.Column(db.String(255), nullable=False)
-#     added = db.Column(db.DateTime(), default=datetime.utcnow, nullable=False)
-#     updated = db.Column(db.DateTime(), default=datetime.utcnow, nullable=False)
-#     destroyed = db.Column(db.DateTime(), nullable=True)
-#     pool_name = db.Column(db.String(80), unique=True, nullable=False)
-#     api_key = db.Column(db.String(80), nullable=False)
+Base = automap_base()
 
 
 def upgrade():
     with op.batch_alter_table('pool', schema=None) as batch_op:
         batch_op.add_column(sa.Column('api_key', sa.String(length=80), nullable=True, unique=False))
-    session = Session(bind=op.get_bind())
+    bind = op.get_bind()
+    Base.prepare(autoload_with=bind)
+    Pool = Base.classes.pool
+    session = Session(bind=bind)
     for pool in session.query(Pool).all():
         pool.api_key = secrets.token_urlsafe(nbytes=32)
     session.commit()