Add soft delete session support

Author: Sakura0122

sqlmodel/__init__.py

@@ -118,6 +118,8 @@
 from .main import Relationship as Relationship
 from .main import SQLModel as SQLModel
 from .orm.session import Session as Session
+from .soft_delete import SoftDeleteMixin as SoftDeleteMixin
+from .soft_delete import SoftDeleteSession as SoftDeleteSession
 from .sql.expression import all_ as all_
 from .sql.expression import and_ as and_
 from .sql.expression import any_ as any_

sqlmodel/soft_delete.py

@@ -0,0 +1,115 @@
+from collections.abc import Mapping, Sequence
+from datetime import datetime, timezone
+from typing import Any, ClassVar, TypeVar, overload
+
+from sqlalchemy import util
+from sqlalchemy.engine.cursor import CursorResult
+from sqlalchemy.engine.result import ScalarResult, TupleResult
+from sqlalchemy.sql.dml import UpdateBase
+
+from .main import Field
+from .orm.session import Session
+from .sql.base import Executable
+from .sql.expression import Select, SelectOfScalar
+
+_TSelectParam = TypeVar("_TSelectParam", bound=Any)
+
+
+class SoftDeleteMixin:
+    __soft_delete_field__: ClassVar[str] = "deleted_at"
+
+    deleted_at: datetime | None = Field(default=None)
+
+
+def get_soft_delete_field(model: type[Any]) -> str | None:
+    field_name = getattr(model, "__soft_delete_field__", None)
+    if not isinstance(field_name, str):
+        return None
+    if not hasattr(model, field_name):
+        return None
+    return field_name
+
+
+class SoftDeleteSession(Session):
+    def delete(self, instance: Any, hard_delete: bool = False) -> None:
+        field_name = get_soft_delete_field(type(instance))
+        if not hard_delete and field_name is not None:
+            setattr(instance, field_name, datetime.now(timezone.utc))
+            self.add(instance)
+            return None
+        return super().delete(instance)
+
+    @overload
+    def exec(
+        self,
+        statement: Select[_TSelectParam],
+        *,
+        params: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None,
+        execution_options: Mapping[str, Any] = util.EMPTY_DICT,
+        bind_arguments: dict[str, Any] | None = None,
+        _parent_execute_state: Any | None = None,
+        _add_event: Any | None = None,
+    ) -> TupleResult[_TSelectParam]: ...
+
+    @overload
+    def exec(
+        self,
+        statement: SelectOfScalar[_TSelectParam],
+        *,
+        params: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None,
+        execution_options: Mapping[str, Any] = util.EMPTY_DICT,
+        bind_arguments: dict[str, Any] | None = None,
+        _parent_execute_state: Any | None = None,
+        _add_event: Any | None = None,
+    ) -> ScalarResult[_TSelectParam]: ...
+
+    @overload
+    def exec(
+        self,
+        statement: UpdateBase,
+        *,
+        params: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None,
+        execution_options: Mapping[str, Any] = util.EMPTY_DICT,
+        bind_arguments: dict[str, Any] | None = None,
+        _parent_execute_state: Any | None = None,
+        _add_event: Any | None = None,
+    ) -> CursorResult[Any]: ...
+
+    def exec(
+        self,
+        statement: Select[_TSelectParam]
+        | SelectOfScalar[_TSelectParam]
+        | Executable[_TSelectParam]
+        | UpdateBase,
+        *,
+        params: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None,
+        execution_options: Mapping[str, Any] = util.EMPTY_DICT,
+        bind_arguments: dict[str, Any] | None = None,
+        _parent_execute_state: Any | None = None,
+        _add_event: Any | None = None,
+    ) -> TupleResult[_TSelectParam] | ScalarResult[_TSelectParam] | CursorResult[Any]:
+        statement = self._filter_soft_deleted(statement)
+        return super().exec(
+            statement,
+            params=params,
+            execution_options=execution_options,
+            bind_arguments=bind_arguments,
+            _parent_execute_state=_parent_execute_state,
+            _add_event=_add_event,
+        )
+
+    def _filter_soft_deleted(self, statement: Any) -> Any:
+        if statement.get_execution_options().get("include_deleted"):
+            return statement
+
+        models: set[type[Any]] = set()
+        for description in getattr(statement, "column_descriptions", ()):
+            model = description.get("entity")
+            if isinstance(model, type):
+                models.add(model)
+
+        for model in models:
+            field_name = get_soft_delete_field(model)
+            if field_name is not None:
+                statement = statement.where(getattr(model, field_name).is_(None))
+        return statement

tests/test_soft_delete.py

@@ -0,0 +1,82 @@
+from datetime import datetime
+
+from sqlmodel import (
+    Field,
+    SoftDeleteMixin,
+    SoftDeleteSession,
+    SQLModel,
+    create_engine,
+    select,
+)
+
+
+def test_soft_delete_filters_select_results() -> None:
+    class Hero(SQLModel, SoftDeleteMixin, table=True):
+        id: int | None = Field(default=None, primary_key=True)
+        name: str
+
+    engine = create_engine("sqlite://")
+    SQLModel.metadata.create_all(engine)
+
+    with SoftDeleteSession(engine) as session:
+        active_hero = Hero(name="Deadpond")
+        deleted_hero = Hero(name="Spider-Boy")
+        session.add(active_hero)
+        session.add(deleted_hero)
+        session.commit()
+
+        session.delete(deleted_hero)
+        session.commit()
+
+        heroes = session.exec(select(Hero)).all()
+        assert heroes == [active_hero]
+        assert isinstance(deleted_hero.deleted_at, datetime)
+
+
+def test_soft_delete_can_include_deleted_results() -> None:
+    class Hero(SQLModel, SoftDeleteMixin, table=True):
+        id: int | None = Field(default=None, primary_key=True)
+        name: str
+
+    engine = create_engine("sqlite://")
+    SQLModel.metadata.create_all(engine)
+
+    with SoftDeleteSession(engine) as session:
+        active_hero = Hero(name="Deadpond")
+        deleted_hero = Hero(name="Spider-Boy")
+        session.add(active_hero)
+        session.add(deleted_hero)
+        session.commit()
+
+        session.delete(deleted_hero)
+        session.commit()
+
+        statement = select(Hero).execution_options(include_deleted=True)
+        heroes = session.exec(statement).all()
+        assert heroes == [active_hero, deleted_hero]
+
+
+def test_soft_delete_supports_custom_field_name() -> None:
+    class Hero(SQLModel, table=True):
+        __soft_delete_field__ = "removed_at"
+
+        id: int | None = Field(default=None, primary_key=True)
+        name: str
+        removed_at: datetime | None = Field(default=None)
+
+    engine = create_engine("sqlite://")
+    SQLModel.metadata.create_all(engine)
+
+    with SoftDeleteSession(engine) as session:
+        active_hero = Hero(name="Deadpond")
+        deleted_hero = Hero(name="Spider-Boy")
+        session.add(active_hero)
+        session.add(deleted_hero)
+        session.commit()
+
+        session.delete(deleted_hero)
+        session.commit()
+
+        heroes = session.exec(select(Hero)).all()
+        assert heroes == [active_hero]
+        assert isinstance(deleted_hero.removed_at, datetime)