feat: support custom executor.

Author: amisadmin

sqlalchemy_database/_abc_async_database.py

@@ -12,15 +12,13 @@
     _P = ParamSpec("_P")
     _R = TypeVar("_R")
 
-
     async def to_thread(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R:  # noqa: E303
         loop = asyncio.get_running_loop()
         ctx = contextvars.copy_context()
         func_call = functools.partial(ctx.run, func, *args, **kwargs)
         return await loop.run_in_executor(None, func_call)
 
-
-class AbcAsyncDatabase(metaclass=abc.ABCMeta):
+class AbcAsyncDatabase(metaclass = abc.ABCMeta):
 
     def __init__(self) -> None:
         for func_name in ['execute', 'scalar', 'scalars_all', 'get', 'delete', 'save', 'run_sync']:

sqlalchemy_database/_abc_async_database.pyi

@@ -2,6 +2,7 @@ import abc
 from typing import Any, Optional, Mapping, Union, Sequence, Type, List, TypeVar, Callable
 
 from sqlalchemy.engine import Result, Connection
+from sqlalchemy.ext.asyncio import AsyncSession, AsyncConnection
 from sqlalchemy.orm import Session
 from sqlalchemy.sql import Executable
 from typing_extensions import ParamSpec, Concatenate
@@ -13,66 +14,75 @@ _R = TypeVar("_R")
 _ExecuteParams = Union[Mapping[Any, Any], Sequence[Mapping[Any, Any]]]
 _ExecuteOptions = Mapping[Any, Any]
 
-
-class AbcAsyncDatabase(metaclass=abc.ABCMeta):
+class AbcAsyncDatabase(metaclass = abc.ABCMeta):
     """`sqlalchemy` asynchronous database abstract base class, not directly instantiated
 
     """
 
     async def async_execute(
-            self,
-            statement: Executable,
-            params: Optional[_ExecuteParams] = None,
-            *,
-            execution_options: Optional[_ExecuteOptions] = None,
-            bind_arguments: Optional[Mapping[str, Any]] = None,
-            commit: bool = True,
-            on_close_pre: Callable[[Result], _T] = None,
-            is_session: bool = True,
-            **kw: Any,
+        self,
+        statement: Executable,
+        params: Optional[_ExecuteParams] = None,
+        *,
+        execution_options: Optional[_ExecuteOptions] = None,
+        bind_arguments: Optional[Mapping[str, Any]] = None,
+        commit: bool = True,
+        on_close_pre: Callable[[Result], _T] = None,
+        is_session: bool = True,
+        executor: Union[Session, Connection, AsyncSession, AsyncConnection, None] = None,
+        **kw: Any,
     ) -> Union[Result, _T]: ...
 
     async def async_scalar(
-            self,
-            statement: Executable,
-            params: Optional[_ExecuteParams] = None,
-            *,
-            execution_options: Optional[_ExecuteOptions] = None,
-            bind_arguments: Optional[Mapping[str, Any]] = None,
-            **kw: Any,
+        self,
+        statement: Executable,
+        params: Optional[_ExecuteParams] = None,
+        *,
+        execution_options: Optional[_ExecuteOptions] = None,
+        bind_arguments: Optional[Mapping[str, Any]] = None,
+        session: Union[Session, AsyncSession, None] = None,
+        **kw: Any,
     ) -> Any: ...
 
     async def async_scalars_all(
-            self,
-            statement: Executable,
-            params: Optional[_ExecuteParams] = None,
-            *,
-            execution_options: Optional[_ExecuteOptions] = None,
-            **kw: Any,
+        self,
+        statement: Executable,
+        params: Optional[_ExecuteParams] = None,
+        *,
+        execution_options: Optional[_ExecuteOptions] = None,
+        session: Union[Session, AsyncSession, None] = None,
+        **kw: Any,
     ) -> List[Any]: ...
 
     async def async_get(
-            self,
-            entity: Type[_T],
-            ident: Any,
-            *,
-            options: Optional[Sequence[Any]] = None,
-            populate_existing: bool = False,
-            with_for_update: Optional[Any] = None,
-            identity_token: Optional[Any] = None,
-            execution_options: Optional[_ExecuteOptions] = None,
+        self,
+        entity: Type[_T],
+        ident: Any,
+        *,
+        options: Optional[Sequence[Any]] = None,
+        populate_existing: bool = False,
+        with_for_update: Optional[Any] = None,
+        identity_token: Optional[Any] = None,
+        execution_options: Optional[_ExecuteOptions] = None,
+        session: Union[Session, AsyncSession, None] = None
     ) -> Optional[_T]: ...
 
     async def async_delete(self, instance: Any) -> None: ...
 
-    async def async_save(self, *instances: Any, refresh: bool = False) -> None: ...
+    async def async_save(
+        self,
+        *instances: Any,
+        refresh: bool = False,
+        session: Union[Session, AsyncSession, None] = None
+    ) -> None: ...
 
     async def async_run_sync(
-            self,
-            fn: Callable[[Concatenate[Union[Session, Connection], _P]], _T],
-            *args: _P.args,
-            commit: bool = True,
-            on_close_pre: Callable[[_T], _R] = None,
-            is_session: bool = True,
-            **kwargs: _P.kwargs
+        self,
+        fn: Callable[[Concatenate[Union[Session, Connection], _P]], _T],
+        *args: _P.args,
+        commit: bool = True,
+        on_close_pre: Callable[[_T], _R] = None,
+        is_session: bool = True,
+        executor: Union[Session, Connection, AsyncSession, AsyncConnection, None] = None,
+        **kwargs: _P.kwargs
     ) -> Union[_T, _R]: ...