test

alembic/env.py

@@ -1,10 +1,14 @@
+import asyncio
+import backend.config as settings
 from logging.config import fileConfig
 
-from sqlalchemy import engine_from_config
-from sqlalchemy import pool
+from sqlalchemy.engine import Connection
 
 from alembic import context
 
+from backend.session import engine
+from models import BaseModel
+
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 config = context.config
@@ -18,12 +22,28 @@ if config.config_file_name is not None:
 # for 'autogenerate' support
 # from myapp import mymodel
 # target_metadata = mymodel.Base.metadata
-target_metadata = None
+target_metadata = BaseModel.metadata
+print(target_metadata.schema)
+
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
 # my_important_option = config.get_main_option("my_important_option")
 # ... etc.
+def include_object(object, name, type_, reflected, compare_to):
+    if type_ == 'table' and object.schema != target_metadata.schema:
+        return False
+    return True
+
+
+def get_url():
+    url = config.get_main_option("sqlalchemy.url").format(
+        PG_LOGIN=settings.PG_LOGIN,
+        PG_PASSWORD=settings.PG_PASSWORD,
+        PG_HOST=settings.PG_HOST,
+        PG_DATABASE=settings.PG_DATABASE,
+    )
+    return url
 
 
 def run_migrations_offline() -> None:
@@ -38,38 +58,59 @@ def run_migrations_offline() -> None:
     script output.
 
     """
-    url = config.get_main_option("sqlalchemy.url")
+    url = get_url()
     context.configure(
         url=url,
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
+        version_table_schema=target_metadata.schema,
+        include_schemas=True
     )
 
     with context.begin_transaction():
+        """
+        By default search_path is setted to "$user",public 
+        that why alembic can't create foreign keys correctly
+        """
+        context.execute('SET search_path TO public')
         context.run_migrations()
 
 
-def run_migrations_online() -> None:
-    """Run migrations in 'online' mode.
+def do_run_migrations(connection: Connection) -> None:
+    # Here we set our custom schema configuration
+    context.configure(
+        connection=connection,
+        target_metadata=target_metadata,
+        version_table="asia_alembic_version",
+        # Table in the db which will save the current alembic version of this schema
+        version_table_schema=target_metadata.schema,  # Alternative: "public", to put version tables in seperate schema
+        include_schemas=True,
+        include_object=include_object,
+    )
+    with context.begin_transaction():
+        # context.execute(f'set search_path to {target_metadata.schema}')
 
-    In this scenario we need to create an Engine
+        context.run_migrations()
+
+
+async def run_async_migrations() -> None:
+    """In this scenario we need to create an Engine
     and associate a connection with the context.
 
     """
-    connectable = engine_from_config(
-        config.get_section(config.config_ini_section, {}),
-        prefix="sqlalchemy.",
-        poolclass=pool.NullPool,
-    )
+    connectable = engine
 
-    with connectable.connect() as connection:
-        context.configure(
-            connection=connection, target_metadata=target_metadata
-        )
+    async with connectable.connect() as connection:
+        await connection.run_sync(do_run_migrations)
 
-        with context.begin_transaction():
-            context.run_migrations()
+    await connectable.dispose()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode."""
+
+    asyncio.run(run_async_migrations())
 
 
 if context.is_offline_mode():