From 181da63c89184376381db2ddf2c4c666073d326f Mon Sep 17 00:00:00 2001
From: Fabio Manganiello <fabio@manganiello.tech>
Date: Sat, 19 Aug 2023 13:02:05 +0200
Subject: [PATCH] Pass the database engine to the Alembic process as an extra
 argument.

If the path of the default database engine is overridden via `--workdir`
option then it won't be visible to the new `python` subprocess spawned
for Alembic.
---
 platypush/entities/_base.py         | 27 ++++++++++++++++++++++-----
 platypush/migrations/alembic/env.py | 13 +++++++------
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/platypush/entities/_base.py b/platypush/entities/_base.py
index c6d5c18b..c3089dde 100644
--- a/platypush/entities/_base.py
+++ b/platypush/entities/_base.py
@@ -24,6 +24,7 @@ from sqlalchemy import (
     UniqueConstraint,
     inspect as schema_inspect,
 )
+from sqlalchemy.engine import Engine
 from sqlalchemy.orm import ColumnProperty, backref, relationship
 from sqlalchemy.orm.exc import ObjectDeletedError
 
@@ -303,6 +304,24 @@ def _discover_entity_types():
                 entities_registry[obj] = {}  # type: ignore
 
 
+def _get_db():
+    """
+    Utility method to get the db plugin.
+    """
+    from platypush.context import get_plugin
+
+    db = get_plugin('db')
+    assert db
+    return db
+
+
+def _get_db_engine() -> Engine:
+    """
+    Utility method to get the db engine.
+    """
+    return _get_db().get_engine()
+
+
 def get_entities_registry() -> EntityRegistryType:
     """
     :returns: A copy of the entities registry.
@@ -314,13 +333,9 @@ def init_entities_db():
     """
     Initializes the entities database.
     """
-    from platypush.context import get_plugin
-
     run_db_migrations()
     _discover_entity_types()
-    db = get_plugin('db')
-    assert db
-    db.create_all(db.get_engine(), Base)
+    _get_db().create_all(_get_db_engine(), Base)
 
 
 def run_db_migrations():
@@ -339,6 +354,8 @@ def run_db_migrations():
             'alembic',
             '-c',
             alembic_ini,
+            '-x',
+            f'DBNAME={_get_db_engine().url}',
             'upgrade',
             'head',
         ],
diff --git a/platypush/migrations/alembic/env.py b/platypush/migrations/alembic/env.py
index a9de8c4f..406b1f6e 100644
--- a/platypush/migrations/alembic/env.py
+++ b/platypush/migrations/alembic/env.py
@@ -74,14 +74,15 @@ def run_migrations_online() -> None:
 
 
 def set_db_engine():
-    db_conf = Config.get('db')
-    assert db_conf, 'Could not retrieve the database configuration'
-    engine = db_conf['engine']
-    assert engine, 'No database engine configured'
+    engine_url = context.get_x_argument(as_dictionary=True).get('DBNAME')
+    if not engine_url:
+        db_conf = Config.get('db')
+        assert db_conf, 'Could not retrieve the database configuration'
+        engine_url = db_conf['engine']
+        assert engine_url, 'No database engine configured'
 
-    config = context.config
     section = config.config_ini_section
-    config.set_section_option(section, 'DB_ENGINE', engine)
+    config.set_section_option(section, 'DB_ENGINE', engine_url)
 
 
 set_db_engine()