From 07ce95e2963ae0e981bf8e74b1d02376899a4314 Mon Sep 17 00:00:00 2001
From: Daniel Grams <grams.daniel@gmail.com>
Date: Sun, 16 Apr 2023 18:40:27 +0200
Subject: [PATCH] Library updates #432

---
 migrations/README      |  2 +-
 migrations/alembic.ini |  7 ++++++-
 migrations/env.py      | 46 ++++++++++++++++++++++++++----------------
 3 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/migrations/README b/migrations/README
index 98e4f9c..0e04844 100644
--- a/migrations/README
+++ b/migrations/README
@@ -1 +1 @@
-Generic single-database configuration.
\ No newline at end of file
+Single-database configuration for Flask.
diff --git a/migrations/alembic.ini b/migrations/alembic.ini
index 6969e2d..ed8a140 100644
--- a/migrations/alembic.ini
+++ b/migrations/alembic.ini
@@ -11,7 +11,7 @@
 
 # Logging configuration
 [loggers]
-keys = root,sqlalchemy,alembic
+keys = root,sqlalchemy,alembic,flask_migrate
 
 [handlers]
 keys = console
@@ -34,6 +34,11 @@ level = INFO
 handlers =
 qualname = alembic
 
+[logger_flask_migrate]
+level = INFO
+handlers =
+qualname = flask_migrate
+
 [handler_console]
 class = StreamHandler
 args = (sys.stderr,)
diff --git a/migrations/env.py b/migrations/env.py
index 22c2ca7..e6176de 100644
--- a/migrations/env.py
+++ b/migrations/env.py
@@ -1,10 +1,8 @@
-from __future__ import with_statement
-
 import logging
 from logging.config import fileConfig
 
 from alembic import context
-from sqlalchemy import engine_from_config, pool
+from flask import current_app
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
@@ -15,17 +13,29 @@ config = context.config
 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
+
+
+def get_engine_url():
+    try:
+        return get_engine().url.render_as_string(hide_password=False).replace("%", "%%")
+    except AttributeError:
+        return str(get_engine().url).replace("%", "%%")
+
+
 # add your model's MetaData object here
 # for 'autogenerate' support
 # from myapp import mymodel
 # target_metadata = mymodel.Base.metadata
-from flask import current_app
-
-config.set_main_option(
-    "sqlalchemy.url",
-    str(current_app.extensions["migrate"].db.engine.url).replace("%", "%%"),
-)
-target_metadata = current_app.extensions["migrate"].db.metadata
+config.set_main_option("sqlalchemy.url", get_engine_url())
+target_db = current_app.extensions["migrate"].db
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
@@ -33,6 +43,12 @@ target_metadata = current_app.extensions["migrate"].db.metadata
 # ... etc.
 
 
+def get_metadata():
+    if hasattr(target_db, "metadatas"):
+        return target_db.metadatas[None]
+    return target_db.metadata
+
+
 def exclude_tables_from_config(config_):
     tables_ = config_.get("tables", None)
     if tables_ is not None:
@@ -65,7 +81,7 @@ def run_migrations_offline():
     url = config.get_main_option("sqlalchemy.url")
     context.configure(
         url=url,
-        target_metadata=target_metadata,
+        target_metadata=get_metadata(),
         literal_binds=True,
         include_object=include_object,
     )
@@ -92,16 +108,12 @@ def run_migrations_online():
                 directives[:] = []
                 logger.info("No changes in schema detected.")
 
-    connectable = engine_from_config(
-        config.get_section(config.config_ini_section),
-        prefix="sqlalchemy.",
-        poolclass=pool.NullPool,
-    )
+    connectable = get_engine()
 
     with connectable.connect() as connection:
         context.configure(
             connection=connection,
-            target_metadata=target_metadata,
+            target_metadata=get_metadata(),
             process_revision_directives=process_revision_directives,
             include_object=include_object,
             **current_app.extensions["migrate"].configure_args