feat: try importing Session from SQLModel

Author: amisadmin

pyproject.toml

@@ -49,6 +49,7 @@ test = [
     "pytest-cov",
     "pytest-asyncio>=0.17",
     "httpx",
+    "sqlmodel",
 ]
 docs = [
     "mkdocs-material>=8.3.8",

sqlalchemy_database/_abc_async_database.pyi

@@ -2,11 +2,17 @@ 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.ext.asyncio import AsyncConnection
 from sqlalchemy.sql import Executable
 from typing_extensions import ParamSpec, Concatenate
 
+try:
+    from sqlmodel import Session
+    from sqlmodel.ext.asyncio.session import AsyncSession
+except ImportError:
+    from sqlalchemy.orm import Session
+    from sqlalchemy.ext.asyncio import AsyncSession
+
 _P = ParamSpec("_P")
 _T = TypeVar("_T")
 _R = TypeVar("_R")

sqlalchemy_database/database.py

@@ -1,12 +1,19 @@
 from typing import Generator, Any, AsyncGenerator, Optional, Mapping, Union, Sequence, Type, List, Callable, TypeVar
 
 from sqlalchemy.engine import Result, Connection
-from sqlalchemy.ext.asyncio import AsyncSession, AsyncEngine, create_async_engine, AsyncConnection
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine, AsyncConnection
 from sqlalchemy.future import Engine, create_engine
-from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.orm import sessionmaker
 from sqlalchemy.sql import Executable, Select
 from typing_extensions import Concatenate, ParamSpec
 
+try:
+    from sqlmodel import Session
+    from sqlmodel.ext.asyncio.session import AsyncSession
+except ImportError:
+    from sqlalchemy.orm import Session
+    from sqlalchemy.ext.asyncio import AsyncSession
+
 from sqlalchemy_database._abc_async_database import AbcAsyncDatabase
 
 _P = ParamSpec("_P")
@@ -368,6 +375,7 @@ class Database(AbcAsyncDatabase):
     def __init__(self, engine: Engine, **session_options):
         super().__init__()
         self.engine: Engine = engine
+        session_options.setdefault('class_', Session)
         self.session_maker: Callable[..., Session] = sessionmaker(self.engine, **session_options)
 
     @classmethod

tests/test_AsyncDatabase.py

@@ -154,3 +154,11 @@ async def test_executor(fake_users):
         users = await db.scalars_all(select(User), session = session)
         for user in users:
             assert user.group is None if user.group_id is None else user.group
+
+async def test_sqlmodel_session(fake_users):
+    from sqlmodel import select
+
+    async with db.session_maker() as session:
+        result = await session.exec(select(User))
+        user = result.first()
+        assert user.id == 1

tests/test_Database.py

@@ -153,3 +153,10 @@ def test_executor(fake_users):
         users = db.scalars_all(select(User), session = session)
         for user in users:
             assert user.group is None if user.group_id is None else user.group
+
+def test_sqlmodel_session(fake_users):
+    from sqlmodel import select
+
+    with db.session_maker() as session:
+        user = session.exec(select(User)).first()
+        assert user.id == 1