Merge pull request #47 from febus982/pydantic_v2_part2

Further pydantic v2 support

README.md

@@ -72,7 +72,7 @@ Example:
 bind = sa_manager.get_bind()
 
 
-class MyModel(bind.model_declarative_base):
+class MyModel(bind.declarative_base):
     pass
 
 
@@ -144,7 +144,7 @@ The `SQLAlchemyRepository` and `SQLAlchemyAsyncRepository` class can be used dir
 from sqlalchemy_bind_manager.repository import SQLAlchemyRepository
 
 
-class MyModel(model_declarative_base):
+class MyModel(declarative_base):
     pass
 
 # Direct usage

docs/index.md

@@ -61,7 +61,7 @@ config = SQLAlchemyConfig(
 sa_manager = SQLAlchemyBindManager(config)
 
 # Declare a model
-class MyModel(sa_manager.get_bind().model_declarative_base):
+class MyModel(sa_manager.get_bind().declarative_base):
     id: Mapped[int] = mapped_column(primary_key=True)
     name: Mapped[str] = mapped_column(String(30))
 

docs/manager/alembic/env.py

@@ -27,7 +27,7 @@
 
 sa_manager = SQLAlchemyBindManager(config=bind_config)
 
-class BookModel(sa_manager.get_bind().model_declarative_base):
+class BookModel(sa_manager.get_bind().declarative_base):
     id = Column(Integer)
     title = Column(String)
 ################################################################

docs/manager/config.md

@@ -23,7 +23,7 @@ Once the bind manager is initialised we can retrieve and use the SQLAlchemyBind
 The `SQLAlchemyBind` class has the following attributes:
 
 * `engine`: The initialised SQLALchemy `Engine`
-* `model_declarative_base`: A base class that can be used to create [declarative models](https://docs.sqlalchemy.org/en/14/orm/mapping_styles.html#declarative-mapping)
+* `declarative_base`: A base class that can be used to create [declarative models](https://docs.sqlalchemy.org/en/14/orm/mapping_styles.html#declarative-mapping)
 * `registry_mapper`: The `registry` associated with the `engine`. It can be used with Alembic or to setup [imperative mapping](https://docs.sqlalchemy.org/en/14/orm/mapping_styles.html#imperative-mapping)
 * `session_class`: The class built by [sessionmaker()](https://docs.sqlalchemy.org/en/14/orm/session_api.html#sqlalchemy.orm.sessionmaker), either `Session` or `AsyncSession`
 

docs/manager/models.md

@@ -10,7 +10,7 @@ from sqlalchemy import String
 
 bind = sa_manager.get_bind()
 
-class MyModel(bind.model_declarative_base):
+class MyModel(bind.declarative_base):
     id: Mapped[int] = mapped_column(primary_key=True)
     name: Mapped[str] = mapped_column(String(30))
 ```

docs/repository/usage.md

@@ -6,7 +6,7 @@ The `SQLAlchemyRepository` and `SQLAlchemyAsyncRepository` class can be used dir
 from sqlalchemy_bind_manager.repository import SQLAlchemyRepository
 
 
-class MyModel(model_declarative_base):
+class MyModel(declarative_base):
     pass
 
 # Direct usage

pyproject.toml

@@ -39,7 +39,7 @@ build-backend = "poetry_dynamic_versioning.backend"
 
 [tool.poetry.dependencies]
 python = ">=3.8,<3.12"
-pydantic = ">=1.10.2, <3.0.0"
+pydantic = "^2.1.1"
 SQLAlchemy = { version = "~2.0.0", extras = ["asyncio", "mypy"] }
 
 [tool.poetry.group.dev]

sqlalchemy_bind_manager/_bind_manager.py

@@ -1,6 +1,6 @@
 from typing import Mapping, MutableMapping, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 from sqlalchemy import MetaData, create_engine
 from sqlalchemy.engine import Engine
 from sqlalchemy.ext.asyncio import (
@@ -10,7 +10,7 @@
     create_async_engine,
 )
 from sqlalchemy.orm import Session, sessionmaker
-from sqlalchemy.orm.decl_api import registry
+from sqlalchemy.orm.decl_api import DeclarativeMeta, registry
 
 from sqlalchemy_bind_manager.exceptions import (
     InvalidConfig,
@@ -32,22 +32,20 @@ class SQLAlchemyAsyncConfig(BaseModel):
 
 class SQLAlchemyBind(BaseModel):
     engine: Engine
-    model_declarative_base: type
+    declarative_base: DeclarativeMeta
     registry_mapper: registry
     session_class: sessionmaker[Session]
 
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
 class SQLAlchemyAsyncBind(BaseModel):
     engine: AsyncEngine
-    model_declarative_base: type
+    declarative_base: DeclarativeMeta
     registry_mapper: registry
     session_class: async_sessionmaker[AsyncSession]
 
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
 _SQLAlchemyConfig = Union[
@@ -123,7 +121,7 @@ def __build_sync_bind(
                 class_=Session,
                 **session_options,
             ),
-            model_declarative_base=registry_mapper.generate_base(),
+            declarative_base=registry_mapper.generate_base(),
         )
 
     def __build_async_bind(
@@ -141,7 +139,7 @@ def __build_async_bind(
                 bind=engine,
                 **session_options,
             ),
-            model_declarative_base=registry_mapper.generate_base(),
+            declarative_base=registry_mapper.generate_base(),
         )
 
     def get_binds(self) -> Mapping[str, Union[SQLAlchemyBind, SQLAlchemyAsyncBind]]:

sqlalchemy_bind_manager/_repository/common.py

@@ -3,7 +3,6 @@
 from typing import Generic, List, TypeVar, Union
 
 from pydantic import BaseModel, StrictInt, StrictStr
-from pydantic.generics import GenericModel
 from sqlalchemy import asc, desc
 
 MODEL = TypeVar("MODEL")
@@ -19,7 +18,7 @@ class PageInfo(BaseModel):
     has_previous_page: bool
 
 
-class PaginatedResult(GenericModel, Generic[MODEL]):
+class PaginatedResult(BaseModel, Generic[MODEL]):
     items: List[MODEL]
     page_info: PageInfo
 
@@ -38,7 +37,7 @@ class CursorPageInfo(BaseModel):
     end_cursor: Union[CursorReference, None] = None
 
 
-class CursorPaginatedResult(GenericModel, Generic[MODEL]):
+class CursorPaginatedResult(BaseModel, Generic[MODEL]):
     items: List[MODEL]
     page_info: CursorPageInfo
 

tests/conftest.py

@@ -102,7 +102,7 @@ def sa_bind(request, sa_manager):
 
 @pytest.fixture
 async def model_classes(sa_bind) -> Tuple[Type, Type]:
-    class ParentModel(sa_bind.model_declarative_base):
+    class ParentModel(sa_bind.declarative_base):
         __tablename__ = "parent_model"
         # required in order to access columns with server defaults
         # or SQL expression defaults, subsequent to a flush, without
@@ -119,7 +119,7 @@ class ParentModel(sa_bind.model_declarative_base):
             lazy="selectin",
         )
 
-    class ChildModel(sa_bind.model_declarative_base):
+    class ChildModel(sa_bind.declarative_base):
         __tablename__ = "child_model"
         # required in order to access columns with server defaults
         # or SQL expression defaults, subsequent to a flush, without

tests/repository/test_composite_pk.py

@@ -26,7 +26,7 @@ def sa_manager() -> SQLAlchemyBindManager:
 def model_class_composite_pk(sa_manager) -> Type:
     default_bind = sa_manager.get_bind()
 
-    class MyModel(default_bind.model_declarative_base):
+    class MyModel(default_bind.declarative_base):
         __tablename__ = "mymodel"
         # required in order to access columns with server defaults
         # or SQL expression defaults, subsequent to a flush, without

tests/repository/test_cursor_paginated_find.py

@@ -7,7 +7,7 @@
 
 @pytest.fixture
 async def model_class_string_pk(sa_bind):
-    class MyModel(sa_bind.model_declarative_base):
+    class MyModel(sa_bind.declarative_base):
         __tablename__ = "mymodel_string_pk"
 
         model_id = Column(String, primary_key=True)

tests/session_handler/test_session_lifecycle.py

@@ -26,17 +26,39 @@ async def test_session_is_removed_on_cleanup(session_handler_class, sa_bind):
     mocked_remove.assert_called_once()
 
 
-async def test_session_is_removed_on_cleanup_even_if_loop_is_not_running(sa_manager):
+def test_session_is_removed_on_cleanup_even_if_loop_is_not_running(sa_manager):
     # Running the test without a loop will trigger the loop creation
     sh = AsyncSessionHandler(sa_manager.get_bind("async"))
     original_session_remove = sh.scoped_session.remove
+    original_get_event_loop = asyncio.get_event_loop
 
     with patch.object(
         sh.scoped_session,
         "remove",
         wraps=original_session_remove,
     ) as mocked_close, patch(
-        "asyncio.get_event_loop", side_effect=RuntimeError()
+        "asyncio.get_event_loop",
+        wraps=original_get_event_loop,
+    ) as mocked_get_event_loop:
+        # This should trigger the garbage collector and close the session
+        sh = None
+
+    mocked_get_event_loop.assert_called_once()
+    mocked_close.assert_called_once()
+
+
+def test_session_is_removed_on_cleanup_even_if_loop_search_errors_out(sa_manager):
+    # Running the test without a loop will trigger the loop creation
+    sh = AsyncSessionHandler(sa_manager.get_bind("async"))
+    original_session_remove = sh.scoped_session.remove
+
+    with patch.object(
+        sh.scoped_session,
+        "remove",
+        wraps=original_session_remove,
+    ) as mocked_close, patch(
+        "asyncio.get_event_loop",
+        side_effect=RuntimeError(),
     ) as mocked_get_event_loop:
         # This should trigger the garbage collector and close the session
         sh = None