1+ import asyncio
2+ import asyncio
13from logging .config import fileConfig
2-
4+
35from alembic import context
4- from sqlalchemy import engine_from_config , pool
5-
6+ from sqlalchemy import pool
7+ from sqlalchemy .ext .asyncio import async_engine_from_config
8+
69from src .models .base import BaseDbModel
710from src .settings import get_settings
8-
9-
10- # this is the Alembic Config object, which provides
11- # access to the values within the .ini file in use.
11+
1212config = context .config
1313settings = get_settings ()
14-
15- # Interpret the config file for Python logging.
16- # This line sets up loggers basically.
14+ config .set_main_option ('sqlalchemy.url' , str (settings .DB_DSN ))
15+
1716if config .config_file_name is not None :
1817 fileConfig (config .config_file_name )
19-
20- # add your model's MetaData object here
21- # for 'autogenerate' support
22- # from myapp import mymodel
23- # target_metadata = mymodel.Base.metadata
18+
2419target_metadata = BaseDbModel .metadata
25-
26- # other values from the config, defined by the needs of env.py,
27- # can be acquired:
28- # my_important_option = config.get_main_option("my_important_option")
29- # ... etc.
30-
31-
32- def run_migrations_offline ():
33- """Run migrations in 'offline' mode.
34-
35- This configures the context with just a URL
36- and not an Engine, though an Engine is acceptable
37- here as well. By skipping the Engine creation
38- we don't even need a DBAPI to be available.
39-
40- Calls to context.execute() here emit the given string to the
41- script output.
42-
43- """
44- url = config .get_main_option ("sqlalchemy.url" )
45- context .configure (
46- url = url ,
47- target_metadata = target_metadata ,
48- literal_binds = True ,
49- dialect_opts = {"paramstyle" : "named" },
50- )
51-
20+
21+ def run_migrations_offline () -> None :
22+ url = config .get_main_option ('sqlalchemy.url' )
23+ context .configure (url = url , target_metadata = target_metadata ,
24+ literal_binds = True , dialect_opts = {'paramstyle' : 'named' })
5225 with context .begin_transaction ():
5326 context .run_migrations ()
54-
55-
56- def run_migrations_online ():
57- """Run migrations in 'online' mode.
58-
59- In this scenario we need to create an Engine
60- and associate a connection with the context.
61-
62- """
63- configuration = config .get_section (config .config_ini_section )
27+
28+ def do_run_migrations (connection ):
29+ context .configure (connection = connection , target_metadata = target_metadata )
30+ with context .begin_transaction ():
31+ context .run_migrations ()
32+
33+ async def run_async_migrations () -> None :
34+ configuration = config .get_section (config .config_ini_section , {})
6435 configuration ['sqlalchemy.url' ] = str (settings .DB_DSN )
65- connectable = engine_from_config (
66- configuration ,
67- prefix = "sqlalchemy." ,
68- poolclass = pool .NullPool ,
69- )
70-
71- with connectable .connect () as connection :
72- context .configure (connection = connection , target_metadata = target_metadata )
73-
74- with context .begin_transaction ():
75- context .run_migrations ()
76-
77-
36+ connectable = async_engine_from_config (
37+ configuration , prefix = 'sqlalchemy.' , poolclass = pool .NullPool )
38+ async with connectable .connect () as connection :
39+ await connection .run_sync (do_run_migrations )
40+ await connectable .dispose ()
41+
42+ def run_migrations_online () -> None :
43+ asyncio .run (run_async_migrations ())
44+
7845if context .is_offline_mode ():
7946 run_migrations_offline ()
8047else :
81- run_migrations_online ()
48+ run_migrations_online ()
0 commit comments