feat(sql): new mixin and updates to base repository

Adds a new mixin plus additional updates to the repository.

This additionally adds new types with Pydantic and utilities to handle
those updates.

.tool-versions

@@ -1,2 +1,2 @@
-python 3.12.0
+python 3.12.3
 pre-commit 3.4.0

pyproject.toml

@@ -12,6 +12,10 @@ python = "^3.12.0"
 sqlalchemy = "^2.0.29"
 sqlalchemy-utils = "^0.41.1"
 inflection = "^0.5.1"
+wrapt = "^1.16.0"
+orjson = "^3.10.3"
+pydantic = "^2.0"
+uuid = "^1.30"
 
 [tool.poetry.group.dev.dependencies]
 pylint = "^3.1.0"

sanctumlabs_dbkit/exceptions.py

@@ -5,3 +5,8 @@
 
 class ModelNotFoundError(Exception):
     """Error indicating a missing model"""
+
+
+class UnsupportedModelOperationError(Exception):
+    """Error indicating an operation on a model is unsupported"""
+    pass

sanctumlabs_dbkit/sql/alembic.py

@@ -0,0 +1,36 @@
+from typing import Any, Literal, Union
+
+from sanctumlabs_dbkit.sql.types import PydanticModel, PydanticModelList
+
+
+def render_item(
+    type_: str, obj: Any, autogen_context: Any
+) -> Union[str, Literal[False]]:
+    """
+    A custom renderer for the `alembic` migration framework which caters for our custom SQLAlchemy pydantic types.
+
+    These types allow for pydantic models to be serialized and deserialized to/from json. Alembic doesn't generate the
+    correct migrations for these cases so we need to do some hackery and override here.
+
+    To leverage the custom renderer, you need to configure it on your migration context in your alembic `env.py`
+
+    ``python
+    with connectable.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata,
+            compare_type=False,
+            render_item=render_item,
+        )
+    ```
+
+    See https://alembic.sqlalchemy.org/en/latest/autogenerate.html#affecting-the-rendering-of-types-themselves
+    See https://gist.github.com/imankulov/4051b7805ad737ace7d8de3d3f934d6b
+    """
+
+    if type_ == "type" and (
+        isinstance(obj, PydanticModelList) or isinstance(obj, PydanticModel)
+    ):
+        return "sa.JSON()"
+
+    return False

sanctumlabs_dbkit/sql/mixins.py

@@ -5,7 +5,7 @@
 from typing import Optional
 from datetime import datetime, timezone
 from uuid import UUID, uuid4
-from sqlalchemy import DateTime, func
+from sqlalchemy import DateTime, func, BIGINT, Identity
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, declared_attr
 from sqlalchemy.dialects.postgresql import UUID as UUIDType
 import inflection
@@ -85,3 +85,19 @@ class TableNameMixin:
     def __tablename__(self) -> str:
         """Table names are snake case plural, for example shipping_records"""
         return inflection.pluralize(inflection.underscore(self.__name__))  # type: ignore[attr-defined]
+
+
+class BigIntIdentityMixin:
+    """
+    A mixin to provide an auto-incrementing bigint primary key column.
+
+    NOTE: usage of this mixin for primary key column purposes is discouraged and should only be used for special
+          cases (e.g. outbox spooler). The UUIDPrimaryKeyMixin is what should typically be used instead (via the
+          BaseModel class)
+    """
+
+    id: Mapped[Optional[int]] = mapped_column(
+        Identity(start=1, cycle=False), primary_key=True, nullable=False, type_=BIGINT
+    )
+
+    pk: str = "id"

sanctumlabs_dbkit/sql/models.py

@@ -1,6 +1,12 @@
 """
 Contains base database models that can be subclassed to add functionality & attributes for database models in an app
 """
+from datetime import datetime, timezone
+from typing import Optional
+from uuid import UUID, uuid4
+
+from sqlalchemy import LargeBinary
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, synonym
 
 from sanctumlabs_dbkit.sql.mixins import (
     AuditedMixin,
@@ -9,6 +15,7 @@
     TableNameMixin,
     TimestampColumnsMixin,
     UUIDPrimaryKeyMixin,
+    BaseIdentityMixin
 )
 
 
@@ -35,6 +42,38 @@ class BaseModel(UUIDPrimaryKeyMixin, AbstractBaseModel):
 
     __abstract__ = True
 
+class BaseOutboxEvent(Base, BaseIdentityMixin, TableNameMixin):
+    """
+    Base model for outbox events. Projects can choose to add additional table args (e.g. custom index) if
+    needed:
+
+    __table_args__ = (
+        Index(
+            ...
+        ),
+    )
+    """
+
+    __abstract__ = True
+
+    uuid: Mapped[UUID] = mapped_column(unique=True, default=uuid4)
+
+    created: Mapped[datetime] = mapped_column(
+        default=lambda: datetime.now(timezone.utc)
+    )
+    destination: Mapped[str]
+    event_type: Mapped[str]
+    correlation_id: Mapped[str]
+    partition_key: Mapped[str]
+    payload: Mapped[bytes] = mapped_column(type_=LargeBinary)
+    sent_time: Mapped[Optional[datetime]]
+    error_message: Mapped[Optional[str]]
+
+    # mimic AbstractBaseModel to play nicely in the base DAO class
+    @declared_attr
+    def created_at(cls) -> Mapped[datetime]:  # noqa: N805
+        return synonym("created")
+
 
 __all__ = [
     "AbstractBaseModel",
@@ -45,4 +84,5 @@ class BaseModel(UUIDPrimaryKeyMixin, AbstractBaseModel):
     "TableNameMixin",
     "TimestampColumnsMixin",
     "UUIDPrimaryKeyMixin",
+    "BaseOutboxEvent",
 ]

sanctumlabs_dbkit/sql/repository.py

@@ -3,7 +3,7 @@
 """
 
 from datetime import datetime, UTC
-from typing import Generic, Any, Optional, Sequence, Type, TypeVar, cast
+from typing import Generic, Any, Optional, Sequence, Type, TypeVar, cast, TypeGuard
 from sqlalchemy import ColumnElement, Select, select
 
 from sanctumlabs_dbkit.exceptions import ModelNotFoundError
@@ -28,11 +28,29 @@ def __init__(self, model: Type[T], session: Session) -> None:
         self.model = model
         self.session = session
 
+    @staticmethod
+    def _supports_soft_deletion(model: Type[T]) -> TypeGuard[Type[AbstractBaseModel]]:
+        """
+        Indicates if the provided model supports soft deletion (has a 'deleted_at' column). This function
+        takes in an argument due to mypy typeguarding requirements, and is thus static.
+        """
+        return issubclass(model, AbstractBaseModel)
+
+    def create(self, refresh: bool = False, **kwargs: Any) -> T:
+        model_instance = self.model(**kwargs)
+        self.session.add(model_instance)
+
+        if refresh:
+            self.session.flush()
+            self.session.refresh(model_instance)
+
+        return cast(T, model_instance)
+
     def query(self, include_deleted: bool = False) -> Select:
         """Returns a select query with the model including deleted records if the include_deleted is set to True"""
         selectable = select(self.model)
 
-        if not include_deleted:
+        if not include_deleted and self._supports_soft_deletion(self.model):
             selectable = selectable.where(
                 self.model.deleted_at == self.model.not_deleted_value()
             )
@@ -67,7 +85,11 @@ def all(self, include_deleted: bool = False) -> Sequence[T]:
     def delete(self, pk: Any) -> None:
         """Deletes a given record with the given primary key"""
         entity = self.find(pk)
-
+
+        # Cast here as mypy type narrowing doesn't infer the type of entity
+        # correctly
+        entity = cast(AbstractBaseModel, self.find(pk))
+
         if entity:
             entity.deleted_at = datetime.now(UTC)
 

sanctumlabs_dbkit/sql/session.py

@@ -28,7 +28,7 @@ def transaction(self, func: FuncT) -> FuncT:
         Example:
 
         ```python
-        from sanctumlabs_dbkit import SessionLocal
+        from sanctumlabs_dbkit.sql import SessionLocal
 
         session = SessionLocal()
 
@@ -59,8 +59,8 @@ def transaction(func: FuncT) -> FuncT:
     Example:
 
     ```python
-    from sanctumlabs_dbkit import SessionLocal
-    from sanctumlabs_dbkit.session import transaction
+    from sanctumlabs_dbkit.sql import SessionLocal
+    from sanctumlabs_dbkit.sql.session import transaction
 
     class UserService():
         def __init__(session: Session):