feat: Optimize the code and add the get_engine_url method

Author: amisadmin

pyproject.toml

@@ -37,7 +37,6 @@ classifiers = [
 ]
 dependencies = [
     "sqlalchemy",
-    "sqlalchemy2-stubs>=0.0.2a29"
 ]
 
 [project.urls]

sqlalchemy_database/__init__.py

@@ -1,7 +1,6 @@
 __version__ = "0.1.0"
 __url__ = "https://github.com/amisadmin/sqlalchemy_database"
 
-from sqlalchemy_database._abc_async_database import AbcAsyncDatabase
-from sqlalchemy_database.database import AsyncDatabase, Database
+from sqlalchemy_database.database import AbcAsyncDatabase, AsyncDatabase, Database
 
 __all__ = ["AsyncDatabase", "Database", "AbcAsyncDatabase"]

sqlalchemy_database/_abc_async_database.py

@@ -1,8 +1,10 @@
 import abc
 import asyncio
 import functools
-from typing import Callable, TypeVar
+from typing import Callable, Dict, TypeVar, Union
 
+from sqlalchemy.ext.asyncio import AsyncEngine
+from sqlalchemy.future import Engine
 from sqlalchemy.orm import scoped_session
 
 try:
@@ -23,17 +25,19 @@ async def to_thread(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs)
 
 
 class AbcAsyncDatabase(metaclass=abc.ABCMeta):  # noqa: B024
-    def __new__(cls, engine, *args, **kwargs):
+
+    _instances: Dict[str, "AbcAsyncDatabase"] = None
+
+    def __new__(cls, engine: Union[Engine, AsyncEngine], *args, **kwargs):
         """Create a new instance of the database class.Each engine url corresponds to a database instance,
         and if it already exists, it is directly returned, otherwise a new instance is created.
         """
-        if not hasattr(cls, "_instances"):
-            cls._instances = {}
+        cls._instances = cls._instances or {}
         if engine.url not in cls._instances:
             cls._instances[engine.url] = super().__new__(cls)
         return cls._instances[engine.url]
 
-    def __init__(self) -> None:
+    def __init__(self, engine: Union[Engine, AsyncEngine], *args, **kwargs) -> None:
         for func_name in {
             "run_sync",
             "begin",
@@ -98,3 +102,16 @@ async def asgi_dispatch(self, request, call_next):
         async with self.__call__(scope=id(request.scope)):
             request.scope[f"__sqlalchemy_database__:{id(self)}"] = self
             return await call_next(request)
+
+    def attach_middleware(self, app):
+        """Attach the middleware to the ASGI application.
+        Example:
+            ```Python
+            app = FastAPI()
+            db = Database.create("sqlite:///test.db")
+            db.attach_middlewares(app)
+            ```
+        """
+        from starlette.middleware.base import BaseHTTPMiddleware
+
+        app.add_middleware(BaseHTTPMiddleware, dispatch=self.asgi_dispatch)

sqlalchemy_database/_abc_async_database.pyi

@@ -11,7 +11,7 @@ from typing import (
     Union,
 )
 
-from sqlalchemy.engine import Connection, Result
+from sqlalchemy.engine import Connection, Engine, Result
 from sqlalchemy.sql import ClauseElement, Executable
 from sqlmodel.engine.result import ScalarResult
 from typing_extensions import Concatenate, ParamSpec
@@ -22,19 +22,23 @@ try:
     from sqlmodel import Session
     from sqlmodel.ext.asyncio.session import AsyncSession
 except ImportError:
-    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
     from sqlalchemy.orm import Session
 
 _P = ParamSpec("_P")
 _T = TypeVar("_T")
 _R = TypeVar("_R")
 
+async def to_thread(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R: ...
+
 _ExecuteParams = Union[Mapping[Any, Any], Sequence[Mapping[Any, Any]]]
 _ExecuteOptions = Mapping[Any, Any]
 
 class AbcAsyncDatabase(metaclass=abc.ABCMeta):
     """`sqlalchemy` asynchronous database abstract base class, not directly instantiated"""
 
+    engine: Union[Engine, AsyncEngine]
+
     async def async_run_sync(
         self,
         fn: Callable[[Concatenate[Union[Session, Connection], _P]], _T],

sqlalchemy_database/database.py

@@ -11,7 +11,7 @@
     Union,
 )
 
-from sqlalchemy.engine import Connection
+from sqlalchemy.engine import URL, Connection
 from sqlalchemy.ext.asyncio import (
     AsyncEngine,
     async_scoped_session,
@@ -51,7 +51,6 @@ def __init__(
             commit_on_exit: Whether to commit the session when the context manager or session generator exits.
             **session_options: The default `session` initialization parameters
         """
-
         self.engine: AsyncEngine = engine
         """`sqlalchemy` Asynchronous Engine
 
@@ -78,7 +77,7 @@ def __init__(
             f"_session_context_var_{id(self)}", default=None
         )
         self.scoped_session: async_scoped_session = async_scoped_session(self.session_maker, scopefunc=self._session_scope.get)
-        super().__init__()
+        super().__init__(engine)
 
     @property
     def session(self) -> AsyncSession:
@@ -115,7 +114,7 @@ def __call__(self, scope: Any = None):
 
     @classmethod
     def create(
-        cls, url: str, *, commit_on_exit: bool = True, session_options: Mapping[str, Any] = None, **kwargs
+        cls, url: Union[str, URL], *, commit_on_exit: bool = True, session_options: Mapping[str, Any] = None, **kwargs
     ) -> "AsyncDatabase":
         """
         Initialize the client with a database connection string
@@ -207,7 +206,7 @@ def __init__(self, engine: Engine, commit_on_exit: bool = True, **session_option
         self._session_scope: ContextVar[Union[str, Session, None]] = ContextVar(f"_session_context_var_{id(self)}", default=None)
         self.scoped_session: scoped_session = scoped_session(self.session_maker, scopefunc=self._session_scope.get)
         """Returns the Session local instance for the current context or current thread."""
-        super().__init__()
+        super().__init__(engine)
 
     @property
     def session(self) -> Session:
@@ -222,7 +221,7 @@ def __call__(self, scope: Any = None):
 
     @classmethod
     def create(
-        cls, url: str, *, commit_on_exit: bool = True, session_options: Optional[Mapping[str, Any]] = None, **kwargs
+        cls, url: Union[str, URL], *, commit_on_exit: bool = True, session_options: Optional[Mapping[str, Any]] = None, **kwargs
     ) -> "Database":
         kwargs.setdefault("future", True)
         engine = create_engine(url, **kwargs)

sqlalchemy_database/utils.py

@@ -0,0 +1,24 @@
+from typing import Union
+
+from sqlalchemy.engine import URL, make_url
+
+SQLALCHEMY_DRIVER = {
+    "sqlite": {"sync": ["pysqlite", "pysqlcipher"], "async": ["aiosqlite"]},
+    "mysql": {"sync": ["pymysql", "mysqldb", "mysqlconnector", "cymysql", "pyodbc"], "async": ["aiomysql", "asyncmy"]},
+    "mariadb": {"sync": ["pymysql", "mysqldb", "mysqlconnector", "cymysql", "pyodbc"], "async": ["aiomysql", "asyncmy"]},
+    "postgresql": {"sync": ["pg8000", "pyscopg2", "psycopg", "psycopg2cffi"], "async": ["asyncpg"]},
+    "oracle": {"sync": ["cx_oracle", "oracledb"], "async": []},
+    "mssql": {"sync": ["pyodbc", "pymssql"], "async": []},
+}
+
+
+def get_engine_url(url: Union[str, URL], sync: bool = True) -> URL:
+    url: URL = make_url(url)
+    backend_name = url.get_backend_name()
+    driver_name = url.get_driver_name()
+    driver_type = "sync" if sync else "async"
+    if driver_name in SQLALCHEMY_DRIVER.get(backend_name, {}).get(driver_type, []):
+        return url
+    new_driver = SQLALCHEMY_DRIVER[backend_name][driver_type][0]
+    url = url.set(drivername=f"{backend_name}+{new_driver}")
+    return url

tests/test_Database.py

@@ -5,6 +5,7 @@
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
+from sqlalchemy_database import Database
 from tests.conftest import Group, User, sync_db
 
 
@@ -124,3 +125,9 @@ def test_ThreadPoolExecutor():
     done, fail = wait(all_task, return_when=ALL_COMPLETED)  # 等待线程运行完毕
     results = {task.result() for task in done}
     assert len(results) == task_count
+
+
+def test_create():
+    sync_db1 = Database.create("sqlite:///amisadmin.db?check_same_thread=False")
+    sync_db2 = Database.create("sqlite:///amisadmin.db?check_same_thread=False")
+    assert sync_db2 is sync_db1

tests/test_utils.py

@@ -0,0 +1,9 @@
+from sqlalchemy_database.utils import get_engine_url
+from tests.conftest import async_db, sync_db
+
+
+def test_get_engine_url():
+    assert get_engine_url(sync_db.engine.url, sync=True) == sync_db.engine.url
+    assert get_engine_url(sync_db.engine.url, sync=False) == async_db.engine.url
+    assert get_engine_url(async_db.engine.url, sync=True) == sync_db.engine.url.set(drivername="sqlite+pysqlite")
+    assert get_engine_url(async_db.engine.url, sync=False) == async_db.engine.url