diff --git a/src/flask_migrate/templates/aioflask-multidb/env.py b/src/flask_migrate/templates/aioflask-multidb/env.py index bc2981c..2d90ccd 100644 --- a/src/flask_migrate/templates/aioflask-multidb/env.py +++ b/src/flask_migrate/templates/aioflask-multidb/env.py @@ -20,14 +20,22 @@ fileConfig(config.config_file_name) logger = logging.getLogger('alembic.env') + +def get_engine(bind_key=None): + try: + # this works with Flask-SQLAlchemy<3 and Alchemical + return current_app.extensions['migrate'].db.get_engine(bind=bind_key) + except TypeError: + # this works with Flask-SQLAlchemy>=3 + return current_app.extensions['migrate'].db.engines.get(bind_key) + + # add your model's MetaData object here # for 'autogenerate' support # from myapp import mymodel # target_metadata = mymodel.Base.metadata config.set_main_option( - 'sqlalchemy.url', - str(current_app.extensions['migrate'].db.get_engine().url).replace( - '%', '%%')) + 'sqlalchemy.url', str(get_engine().url).replace('%', '%%')) bind_names = [] if current_app.config.get('SQLALCHEMY_BINDS') is not None: bind_names = list(current_app.config['SQLALCHEMY_BINDS'].keys()) @@ -39,8 +47,7 @@ for bind in bind_names: context.config.set_section_option( bind, "sqlalchemy.url", - str(current_app.extensions['migrate'].db.get_engine( - bind=bind).url).replace('%', '%%')) + str(get_engine(bind_key=bind).url).replace('%', '%%')) target_db = current_app.extensions['migrate'].db @@ -166,12 +173,11 @@ async def run_migrations_online(): # for the direct-to-DB use case, start a transaction on all # engines, then run all migrations, then commit all transactions. engines = { - '': {'engine': current_app.extensions['migrate'].db.get_engine()} + '': {'engine': get_engine()} } for name in bind_names: engines[name] = rec = {} - rec['engine'] = current_app.extensions['migrate'].db.get_engine( - bind=name) + rec['engine'] = get_engine(bind_key=name) for name, rec in engines.items(): engine = rec['engine'] diff --git a/src/flask_migrate/templates/aioflask/env.py b/src/flask_migrate/templates/aioflask/env.py index bfb2957..feab4a8 100644 --- a/src/flask_migrate/templates/aioflask/env.py +++ b/src/flask_migrate/templates/aioflask/env.py @@ -17,14 +17,22 @@ fileConfig(config.config_file_name) logger = logging.getLogger('alembic.env') + +def get_engine(): + try: + # this works with Flask-SQLAlchemy<3 and Alchemical + return current_app.extensions['migrate'].db.get_engine() + except TypeError: + # this works with Flask-SQLAlchemy>=3 + return current_app.extensions['migrate'].db.engine + + # add your model's MetaData object here # for 'autogenerate' support # from myapp import mymodel # target_metadata = mymodel.Base.metadata config.set_main_option( - 'sqlalchemy.url', - str(current_app.extensions['migrate'].db.get_engine().url).replace( - '%', '%%')) + 'sqlalchemy.url', str(get_engine().url).replace('%', '%%')) target_db = current_app.extensions['migrate'].db # other values from the config, defined by the needs of env.py, @@ -90,7 +98,7 @@ async def run_migrations_online(): """ - connectable = current_app.extensions['migrate'].db.get_engine() + connectable = get_engine() async with connectable.connect() as connection: await connection.run_sync(do_run_migrations) diff --git a/src/flask_migrate/templates/flask-multidb/env.py b/src/flask_migrate/templates/flask-multidb/env.py index 476b9bf..953e2b2 100644 --- a/src/flask_migrate/templates/flask-multidb/env.py +++ b/src/flask_migrate/templates/flask-multidb/env.py @@ -22,12 +22,11 @@ def get_engine(bind_key=None): try: - # this works with Flask-SQLAlchemy>=3 - return current_app.extensions['migrate'].db.get_engine( - bind_key=bind_key) - except TypeError: - # this works with Flask-SQLAlchemy<3 + # this works with Flask-SQLAlchemy<3 and Alchemical return current_app.extensions['migrate'].db.get_engine(bind=bind_key) + except TypeError: + # this works with Flask-SQLAlchemy>=3 + return current_app.extensions['migrate'].db.engines.get(bind_key) # add your model's MetaData object here @@ -35,9 +34,7 @@ def get_engine(bind_key=None): # from myapp import mymodel # target_metadata = mymodel.Base.metadata config.set_main_option( - 'sqlalchemy.url', - str(current_app.extensions['migrate'].db.get_engine().url).replace( - '%', '%%')) + 'sqlalchemy.url', str(get_engine().url).replace('%', '%%')) bind_names = [] if current_app.config.get('SQLALCHEMY_BINDS') is not None: bind_names = list(current_app.config['SQLALCHEMY_BINDS'].keys()) @@ -138,7 +135,7 @@ def process_revision_directives(context, revision, directives): # for the direct-to-DB use case, start a transaction on all # engines, then run all migrations, then commit all transactions. engines = { - '': {'engine': current_app.extensions['migrate'].db.get_engine()} + '': {'engine': get_engine()} } for name in bind_names: engines[name] = rec = {} diff --git a/src/flask_migrate/templates/flask/env.py b/src/flask_migrate/templates/flask/env.py index 55e9df9..fa412e8 100644 --- a/src/flask_migrate/templates/flask/env.py +++ b/src/flask_migrate/templates/flask/env.py @@ -16,14 +16,22 @@ fileConfig(config.config_file_name) logger = logging.getLogger('alembic.env') + +def get_engine(): + try: + # this works with Flask-SQLAlchemy<3 and Alchemical + return current_app.extensions['migrate'].db.get_engine() + except TypeError: + # this works with Flask-SQLAlchemy>=3 + return current_app.extensions['migrate'].db.engine + + # add your model's MetaData object here # for 'autogenerate' support # from myapp import mymodel # target_metadata = mymodel.Base.metadata config.set_main_option( - 'sqlalchemy.url', - str(current_app.extensions['migrate'].db.get_engine().url).replace( - '%', '%%')) + 'sqlalchemy.url', str(get_engine().url).replace('%', '%%')) target_db = current_app.extensions['migrate'].db # other values from the config, defined by the needs of env.py, @@ -77,7 +85,7 @@ def process_revision_directives(context, revision, directives): directives[:] = [] logger.info('No changes in schema detected.') - connectable = current_app.extensions['migrate'].db.get_engine() + connectable = get_engine() with connectable.connect() as connection: context.configure(