2020
2121from typing import Mapping , MutableMapping , Union
2222
23- from pydantic import BaseModel , ConfigDict
23+ from pydantic import BaseModel , ConfigDict , StrictBool
2424from sqlalchemy import MetaData , create_engine
2525from sqlalchemy .engine import Engine
2626from sqlalchemy .ext .asyncio import (
4040
4141class SQLAlchemyConfig (BaseModel ):
4242 """
43- Configuration for synchronous engines
44- """
45-
46- engine_url : str
47- engine_options : Union [dict , None ] = None
48- session_options : Union [dict , None ] = None
49-
50-
51- class SQLAlchemyAsyncConfig (BaseModel ):
52- """
53- Configuration for asynchronous engines
43+ Configuration for engines
5444 """
5545
5646 engine_url : str
5747 engine_options : Union [dict , None ] = None
5848 session_options : Union [dict , None ] = None
49+ async_engine : StrictBool = False
5950
6051
6152class SQLAlchemyBind (BaseModel ):
@@ -85,9 +76,8 @@ class SQLAlchemyBindManager:
8576 def __init__ (
8677 self ,
8778 config : Union [
88- Mapping [str , Union [ SQLAlchemyConfig , SQLAlchemyAsyncConfig ] ],
79+ Mapping [str , SQLAlchemyConfig ],
8980 SQLAlchemyConfig ,
90- SQLAlchemyAsyncConfig ,
9181 ],
9282 ) -> None :
9383 self .__binds = {}
@@ -97,18 +87,10 @@ def __init__(
9787 else :
9888 self .__init_bind (DEFAULT_BIND_NAME , config )
9989
100- def __init_bind (
101- self , name : str , config : Union [SQLAlchemyConfig , SQLAlchemyAsyncConfig ]
102- ):
103- if not any (
104- [
105- isinstance (config , SQLAlchemyConfig ),
106- isinstance (config , SQLAlchemyAsyncConfig ),
107- ]
108- ):
90+ def __init_bind (self , name : str , config : SQLAlchemyConfig ):
91+ if not isinstance (config , SQLAlchemyConfig ):
10992 raise InvalidConfigError (
110- f"Config for bind `{ name } ` is not a SQLAlchemyConfig"
111- f" or SQLAlchemyAsyncConfig object"
93+ f"Config for bind `{ name } ` is not a SQLAlchemyConfig" f"object"
11294 )
11395
11496 engine_options : dict = config .engine_options or {}
@@ -119,7 +101,7 @@ def __init_bind(
119101 session_options .setdefault ("expire_on_commit" , False )
120102 session_options .setdefault ("autobegin" , False )
121103
122- if isinstance ( config , SQLAlchemyAsyncConfig ) :
104+ if config . async_engine :
123105 self .__binds [name ] = self .__build_async_bind (
124106 engine_url = config .engine_url ,
125107 engine_options = engine_options ,
0 commit comments