feat: Support automatic commit session before closed.

Author: amisadmin

sqlalchemy_database/_abc_async_database.py

@@ -40,6 +40,7 @@ async def asgi_dispatch(self, request, call_next):
             app.add_middleware(BaseHTTPMiddleware, dispatch=db.asgi_dispatch)
             ```
         """
-        async with self.__call__():
-            response = await call_next(request)
-        return response
+        if self.session is None:  # bind session to request
+            async with self.__call__():
+                return await call_next(request)
+        return await call_next(request)

sqlalchemy_database/database.py

@@ -40,11 +40,17 @@
 class AsyncDatabase(AbcAsyncDatabase):
     """`sqlalchemy` asynchronous database client"""
 
-    def __init__(self, engine: AsyncEngine, **session_options):
+    def __init__(
+        self,
+        engine: AsyncEngine,
+        commit_on_close: bool = False,
+        **session_options,
+    ):
         """
         Initialize the client through the asynchronous engine
         Args:
             engine: Asynchronous Engine
+            commit_on_close: Whether to commit the session when the context manager or session generator exits.
             **session_options: The default `session` initialization parameters
         """
         super().__init__()
@@ -57,6 +63,8 @@ def __init__(self, engine: AsyncEngine, **session_options):
                 await conn.run_sync(SQLModel.metadata.create_all)
             ```
         """
+        self.commit_on_close: bool = commit_on_close
+        """Whether to commit the session when the context manager or session generator exits."""
         session_options.setdefault("class_", AsyncSession)
         self.session_maker: Callable[..., AsyncSession] = sessionmaker(self.engine, **session_options)
         """`sqlalchemy` session factory function
@@ -97,11 +105,14 @@ def __call__(self):
         return AsyncSessionContextVarManager(self)
 
     @classmethod
-    def create(cls, url: str, *, session_options: Mapping[str, Any] = None, **kwargs) -> "AsyncDatabase":
+    def create(
+        cls, url: str, *, commit_on_close: bool = False, session_options: Mapping[str, Any] = None, **kwargs
+    ) -> "AsyncDatabase":
         """
         Initialize the client with a database connection string
         Args:
             url: Asynchronous database connection string
+            commit_on_close: Whether to commit the session when the context manager or session generator exits.
             session_options: The default `session` initialization parameters
             **kwargs: Asynchronous engine initialization parameters
 
@@ -111,7 +122,7 @@ def create(cls, url: str, *, session_options: Mapping[str, Any] = None, **kwargs
         kwargs.setdefault("future", True)
         engine = create_async_engine(url, **kwargs)
         session_options = session_options or {}
-        return cls(engine, **session_options)
+        return cls(engine, commit_on_close=commit_on_close, **session_options)
 
     async def session_generator(self) -> AsyncGenerator[AsyncSession, Any]:
         """AsyncSession Generator, available for FastAPI dependencies.
@@ -125,6 +136,8 @@ async def session_generator(self) -> AsyncGenerator[AsyncSession, Any]:
         """
         async with self.session_maker() as session:
             yield session
+            if self.commit_on_close:
+                await session.commit()
 
     async def _executor_maker(
         self, executor: Union[AsyncSession, AsyncConnection, None] = None, is_session: bool = True
@@ -421,9 +434,10 @@ async def refresh(self, instance, attribute_names=None, with_for_update=None, se
 class Database(AbcAsyncDatabase):
     """`sqlalchemy` synchronous database client"""
 
-    def __init__(self, engine: Engine, **session_options):
+    def __init__(self, engine: Engine, commit_on_close: bool = False, **session_options):
         super().__init__()
         self.engine: Engine = engine
+        self.commit_on_close: bool = commit_on_close
         session_options.setdefault("class_", Session)
         self.session_maker: Callable[..., Session] = sessionmaker(self.engine, **session_options)
         self._session_context_var: ContextVar[Optional[Session]] = ContextVar("_session_context_var", default=None)
@@ -437,7 +451,9 @@ def __call__(self):
         return SessionContextVarManager(self)
 
     @classmethod
-    def create(cls, url: str, *, session_options: Optional[Mapping[str, Any]] = None, **kwargs) -> "Database":
+    def create(
+        cls, url: str, *, commit_on_close: bool = False, session_options: Optional[Mapping[str, Any]] = None, **kwargs
+    ) -> "Database":
         kwargs.setdefault("future", True)
         engine = create_engine(url, **kwargs)
         session_options = session_options or {}
@@ -446,6 +462,8 @@ def create(cls, url: str, *, session_options: Optional[Mapping[str, Any]] = None
     def session_generator(self) -> Generator[Session, Any, None]:
         with self.session_maker() as session:
             yield session
+            if self.commit_on_close:
+                session.commit()
 
     def _executor_maker(
         self, executor: Union[Session, Connection, None] = None, is_session: bool = True
@@ -626,6 +644,8 @@ async def __aexit__(self, exc_type, exc_value, traceback):
         session = self.db._session_context_var.get()
         if exc_type is not None:
             await session.rollback()
+        if self.db.commit_on_close:
+            await session.commit()
         await session.close()
         self.db._session_context_var.reset(self.token)
 
@@ -644,6 +664,8 @@ def __exit__(self, exc_type, exc_value, traceback):
         session = self.db._session_context_var.get()
         if exc_type is not None:
             session.rollback()
+        if self.db.commit_on_close:
+            session.commit()
         session.close()
         self.db._session_context_var.reset(self.token)