diff --git a/sqlmodel/__init__.py b/sqlmodel/__init__.py index f7359f0c27..6fd7548543 100644 --- a/sqlmodel/__init__.py +++ b/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_ diff --git a/sqlmodel/soft_delete.py b/sqlmodel/soft_delete.py new file mode 100644 index 0000000000..2f0e8a7f38 --- /dev/null +++ b/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 diff --git a/tests/test_soft_delete.py b/tests/test_soft_delete.py new file mode 100644 index 0000000000..3722dcb6b6 --- /dev/null +++ b/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)