activemodel 0.3.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

activemodel/__init__.py

@@ -1,6 +1,4 @@
 from .base_model import BaseModel
-from .query_wrapper import QueryWrapper
-from .timestamps import TimestampMixin
 
-# TODO need a way to specify the session generator
-# TODO need a way to specify the session generator
+# from .field import Field
+from .session_manager import SessionManager, get_engine, get_session, init

activemodel/base_model.py

@@ -1,12 +1,38 @@
 import json
 import typing as t
+from uuid import UUID
 
 import pydash
 import sqlalchemy as sa
-from sqlalchemy.orm import declared_attr
-from sqlmodel import Session, SQLModel
-
+import sqlmodel as sm
+from sqlalchemy import Connection, event
+from sqlalchemy.orm import Mapper, declared_attr
+from sqlmodel import Field, MetaData, Session, SQLModel, select
+from typeid import TypeID
+
+# NOTE: this patches a core method in sqlmodel to support db comments
+from . import get_column_from_field_patch  # noqa: F401
+from .logger import logger
 from .query_wrapper import QueryWrapper
+from .session_manager import get_session
+
+POSTGRES_INDEXES_NAMING_CONVENTION = {
+    "ix": "%(column_0_label)s_idx",
+    "uq": "%(table_name)s_%(column_0_name)s_key",
+    "ck": "%(table_name)s_%(constraint_name)s_check",
+    "fk": "%(table_name)s_%(column_0_name)s_fkey",
+    "pk": "%(table_name)s_pkey",
+}
+"""
+By default, the foreign key naming convention in sqlalchemy do not create unique identifiers when there are multiple
+foreign keys in a table. This naming convention is a workaround to fix this issue:
+
+- https://github.com/zhanymkanov/fastapi-best-practices?tab=readme-ov-file#set-db-keys-naming-conventions
+- https://github.com/fastapi/sqlmodel/discussions/1213
+- Implementation lifted from: https://github.com/AlexanderZharyuk/billing-service/blob/3c8aaf19ab7546b97cc4db76f60335edec9fc79d/src/models.py#L24
+"""
+
+SQLModel.metadata.naming_convention = POSTGRES_INDEXES_NAMING_CONVENTION
 
 
 class BaseModel(SQLModel):
@@ -14,32 +40,110 @@ class BaseModel(SQLModel):
     Base model class to inherit from so we can hate python less
 
     https://github.com/woofz/sqlmodel-basecrud/blob/main/sqlmodel_basecrud/basecrud.py
+
+    - {before,after} lifecycle hooks are modeled after Rails.
+    - class docstrings are converd to table-level comments
+    - save(), delete(), select(), where(), and other easy methods you would expect
+    - Fixes foreign key naming conventions
     """
 
-    # TODO implement actually calling these hooks
+    # this is used for table-level comments
+    __table_args__ = None
+
+    @classmethod
+    def __init_subclass__(cls, **kwargs):
+        "Setup automatic sqlalchemy lifecycle events for the class"
+
+        super().__init_subclass__(**kwargs)
+
+        from sqlmodel._compat import set_config_value
+
+        # enables field-level docstrings on the pydanatic `description` field, which we then copy into
+        # sa_args, which is persisted to sql table comments
+        set_config_value(model=cls, parameter="use_attribute_docstrings", value=True)
+
+        cls._apply_class_doc()
+
+        def event_wrapper(method_name: str):
+            """
+            This does smart heavy lifting for us to make sqlalchemy lifecycle events nicer to work with:
+
+            * Passes the target first to the lifecycle method, so it feels like an instance method
+            * Allows as little as a single positional argument, so methods can be simple
+            * Removes the need for decorators or anything fancy on the subclass
+            """
+
+            def wrapper(mapper: Mapper, connection: Connection, target: BaseModel):
+                if hasattr(cls, method_name):
+                    method = getattr(cls, method_name)
+
+                    if callable(method):
+                        arg_count = method.__code__.co_argcount
 
-    def before_delete(self):
-        pass
+                        if arg_count == 1:  # Just self/cls
+                            method(target)
+                        elif arg_count == 2:  # Self, mapper
+                            method(target, mapper)
+                        elif arg_count == 3:  # Full signature
+                            method(target, mapper, connection)
+                        else:
+                            raise TypeError(
+                                f"Method {method_name} must accept either 1 to 3 arguments, got {arg_count}"
+                            )
+                    else:
+                        logger.warning(
+                            "SQLModel lifecycle hook found, but not callable hook_name=%s",
+                            method_name,
+                        )
 
-    def after_delete(self):
-        pass
+            return wrapper
 
-    def before_save(self):
-        pass
+        event.listen(cls, "before_insert", event_wrapper("before_insert"))
+        event.listen(cls, "before_update", event_wrapper("before_update"))
 
-    def after_save(self):
-        pass
+        # before_save maps to two type of events
+        event.listen(cls, "before_insert", event_wrapper("before_save"))
+        event.listen(cls, "before_update", event_wrapper("before_save"))
 
-    def before_update(self):
-        pass
+        # now, let's handle after_* variants
+        event.listen(cls, "after_insert", event_wrapper("after_insert"))
+        event.listen(cls, "after_update", event_wrapper("after_update"))
 
-    def after_update(self):
-        pass
+        # after_save maps to two type of events
+        event.listen(cls, "after_insert", event_wrapper("after_save"))
+        event.listen(cls, "after_update", event_wrapper("after_save"))
 
+        # def foreign_key()
+        # table.id
+
+    @classmethod
+    def _apply_class_doc(cls):
+        """
+        Pull class-level docstring into a table comment.
+
+        This will help AI SQL writers like: https://github.com/iloveitaly/sql-ai-prompt-generator
+        """
+
+        doc = cls.__doc__.strip() if cls.__doc__ else None
+
+        if doc:
+            table_args = getattr(cls, "__table_args__", None)
+
+            if table_args is None:
+                cls.__table_args__ = {"comment": doc}
+            elif isinstance(table_args, dict):
+                table_args.setdefault("comment", doc)
+            else:
+                raise ValueError("Unexpected __table_args__ type")
+
+    # TODO no type check decorator here
     @declared_attr
     def __tablename__(cls) -> str:
         """
         Automatically generates the table name for the model by converting the class name from camel case to snake case.
+        This is the recommended format for table names:
+
+        https://wiki.postgresql.org/wiki/Don%27t_Do_This#Don.27t_use_upper_case_table_or_column_names
 
         By default, the class is lower cased which makes it harder to read.
 
@@ -50,74 +154,137 @@ class BaseModel(SQLModel):
         """
         return pydash.strings.snake_case(cls.__name__)
 
+    @classmethod
+    def foreign_key(cls, **kwargs):
+        """
+        Returns a Field object referencing the foreign key of the model.
+        """
+
+        field_options = {"nullable": False} | kwargs
+
+        return Field(
+            # TODO id field is hard coded, should pick the PK field in case it's different
+            sa_type=cls.model_fields["id"].sa_column.type,  # type: ignore
+            foreign_key=f"{cls.__tablename__}.id",
+            **field_options,
+        )
+
     @classmethod
     def select(cls, *args):
+        "create a query wrapper to easily run sqlmodel queries on this model"
         return QueryWrapper[cls](cls, *args)
 
+    def delete(self):
+        with get_session() as session:
+            if old_session := Session.object_session(self):
+                old_session.expunge(self)
+
+            session.delete(self)
+            session.commit()
+            return True
+
     def save(self):
-        old_session = Session.object_session(self)
         with get_session() as session:
-            if old_session:
+            if old_session := Session.object_session(self):
                 # I was running into an issue where the object was already
                 # associated with a session, but the session had been closed,
                 # to get around this, you need to remove it from the old one,
                 # then add it to the new one (below)
                 old_session.expunge(self)
 
-            self.before_update()
-            # self.before_save()
-
             session.add(self)
+            # NOTE very important method! This triggers sqlalchemy lifecycle hooks automatically
             session.commit()
             session.refresh(self)
 
-            self.after_update()
-            # self.after_save()
-
-            return self
+        return self
 
-            # except IntegrityError:
-            #     log.quiet(f"{self} already exists in the database.")
-            #     session.rollback()
+        # except IntegrityError:
+        #     log.quiet(f"{self} already exists in the database.")
+        #     session.rollback()
 
     # TODO shouldn't this be handled by pydantic?
     def json(self, **kwargs):
         return json.dumps(self.dict(), default=str, **kwargs)
 
+    # TODO should move this to the wrapper
     @classmethod
-    def count(cls):
+    def count(cls) -> int:
         """
         Returns the number of records in the database.
         """
-        # TODO should move this to the wrapper
         with get_session() as session:
-            query = sql.select(sql.func.count()).select_from(cls)
-            return session.exec(query).one()
+            return session.scalar(sm.select(sm.func.count()).select_from(cls))
+
+    # TODO got to be a better way to fwd these along...
+    @classmethod
+    def first(cls):
+        return cls.select().first()
+
+    # TODO throw an error if this field is set on the model
+    def is_new(self) -> bool:
+        return not self._sa_instance_state.has_identity
+
+    @classmethod
+    def find_or_create_by(cls, **kwargs):
+        """
+        Find record or create it with the passed args if it doesn't exist.
+        """
+
+        result = cls.get(**kwargs)
+
+        if result:
+            return result
+
+        new_model = cls(**kwargs)
+        new_model.save()
+
+        return new_model
+
+    @classmethod
+    def find_or_initialize_by(cls, **kwargs):
+        """
+        Unfortunately, unlike ruby, python does not have a great lambda story. This makes writing convenience methods
+        like this a bit more difficult.
+        """
+
+        result = cls.get(**kwargs)
+
+        if result:
+            return result
+
+        new_model = cls(**kwargs)
+        return new_model
 
     # TODO what's super dangerous here is you pass a kwarg which does not map to a specific
     #      field it will result in `True`, which will return all records, and not give you any typing
     #      errors. Dangerous when iterating on structure quickly
     # TODO can we pass the generic of the superclass in?
+    # TODO can we type the method signature a bit better?
+    # def get(cls, *args: sa.BinaryExpression, **kwargs: t.Any):
     @classmethod
-    def get(cls, *args: sa.BinaryExpression, **kwargs: t.Any):
+    def get(cls, *args: t.Any, **kwargs: t.Any):
         """
         Gets a single record from the database. Pass an PK ID or a kwarg to filter by.
         """
 
+        # TODO id is hardcoded, not good! Need to dynamically pick the best uid field
+        id_field_name = "id"
+
         # special case for getting by ID
-        if len(args) == 1 and isinstance(args[0], int):
-            # TODO id is hardcoded, not good! Need to dynamically pick the best uid field
-            kwargs["id"] = args[0]
-            args = []
+        if len(args) == 1 and isinstance(args[0], (int, TypeID, str, UUID)):
+            kwargs[id_field_name] = args[0]
+            args = ()
+
+        statement = select(cls).filter(*args).filter_by(**kwargs)
 
-        statement = sql.select(cls).filter(*args).filter_by(**kwargs)
         with get_session() as session:
             return session.exec(statement).first()
 
     @classmethod
     def all(cls):
         with get_session() as session:
-            results = session.exec(sql.select(cls))
+            results = session.exec(sm.select(cls))
 
             # TODO do we need this or can we just return results?
             for result in results:
@@ -131,7 +298,7 @@ class BaseModel(SQLModel):
         Helpful for testing and console debugging.
         """
 
-        query = sql.select(cls).order_by(sql.func.random()).limit(1)
+        query = sm.select(cls).order_by(sa.sql.func.random()).limit(1)
 
         with get_session() as session:
             return session.exec(query).one()

activemodel/celery.py

@@ -0,0 +1,28 @@
+"""
+Do not import unless you have Celery/Kombu installed.
+
+In order for TypeID objects to be properly handled by celery, a custom encoder must be registered.
+"""
+
+from kombu.utils.json import register_type
+from typeid import TypeID
+
+
+def register_celery_typeid_encoder():
+    "this ensures TypeID objects passed as arguments to a delayed function are properly serialized"
+
+    def class_full_name(clz) -> str:
+        return ".".join([clz.__module__, clz.__qualname__])
+
+    def _encoder(obj: TypeID) -> str:
+        return str(obj)
+
+    def _decoder(data: str) -> TypeID:
+        return TypeID.from_string(data)
+
+    register_type(
+        TypeID,
+        class_full_name(TypeID),
+        encoder=_encoder,
+        decoder=_decoder,
+    )

activemodel/errors.py

@@ -0,0 +1,6 @@
+class TypeIDValidationError(ValueError):
+    """
+    Raised when a TypeID is invalid in some way
+    """
+
+    pass

activemodel/get_column_from_field_patch.py

@@ -0,0 +1,137 @@
+"""
+Pydantic has a great DX for adding docstrings to fields. This allows devs to easily document the fields of a model.
+
+Making sure these docstrings make their way to the DB schema is helpful for a bunch of reasons (LLM understanding being one of them).
+
+This patch mutates a core sqlmodel function which translates pydantic FieldInfo objects into sqlalchemy Column objects. It adds the field description as a comment to the column.
+
+Note that FieldInfo *from pydantic* is used when a "bare" field is defined. This can be confusing, because when inspecting model fields, the class name looks exactly the same.
+"""
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Sequence,
+    cast,
+)
+
+import sqlmodel
+from pydantic.fields import FieldInfo as PydanticFieldInfo
+from sqlalchemy import (
+    Column,
+    ForeignKey,
+)
+from sqlmodel._compat import (  # type: ignore[attr-defined]
+    IS_PYDANTIC_V2,
+    ModelMetaclass,
+    Representation,
+    Undefined,
+    UndefinedType,
+    is_field_noneable,
+)
+from sqlmodel.main import FieldInfo, get_sqlalchemy_type
+
+from activemodel.utils import hash_function_code
+
+if TYPE_CHECKING:
+    from pydantic._internal._model_construction import ModelMetaclass as ModelMetaclass
+    from pydantic._internal._repr import Representation as Representation
+    from pydantic_core import PydanticUndefined as Undefined
+    from pydantic_core import PydanticUndefinedType as UndefinedType
+
+
+assert (
+    hash_function_code(sqlmodel.main.get_column_from_field)
+    == "398006ef8fd8da191ca1a271ef25b6e135da0f400a80df2f29526d8674f9ec51"
+)
+
+
+def get_column_from_field(field: PydanticFieldInfo | FieldInfo) -> Column:  # type: ignore
+    """
+    Takes a field definition, which can either come from the sqlmodel FieldInfo class or the pydantic variant of that class,
+    and converts it into a sqlalchemy Column object.
+    """
+    if IS_PYDANTIC_V2:
+        field_info = field
+    else:
+        field_info = field.field_info
+
+    sa_column = getattr(field_info, "sa_column", Undefined)
+    if isinstance(sa_column, Column):
+        # IMPORTANT: change from the original function
+        if not sa_column.comment and (field_comment := field_info.description):
+            sa_column.comment = field_comment
+        return sa_column
+
+    primary_key = getattr(field_info, "primary_key", Undefined)
+    if primary_key is Undefined:
+        primary_key = False
+
+    index = getattr(field_info, "index", Undefined)
+    if index is Undefined:
+        index = False
+
+    nullable = not primary_key and is_field_noneable(field)
+    # Override derived nullability if the nullable property is set explicitly
+    # on the field
+    field_nullable = getattr(field_info, "nullable", Undefined)  # noqa: B009
+    if field_nullable is not Undefined:
+        assert not isinstance(field_nullable, UndefinedType)
+        nullable = field_nullable
+    args = []
+    foreign_key = getattr(field_info, "foreign_key", Undefined)
+    if foreign_key is Undefined:
+        foreign_key = None
+    unique = getattr(field_info, "unique", Undefined)
+    if unique is Undefined:
+        unique = False
+    if foreign_key:
+        if field_info.ondelete == "SET NULL" and not nullable:
+            raise RuntimeError('ondelete="SET NULL" requires nullable=True')
+        assert isinstance(foreign_key, str)
+        ondelete = getattr(field_info, "ondelete", Undefined)
+        if ondelete is Undefined:
+            ondelete = None
+        assert isinstance(ondelete, (str, type(None)))  # for typing
+        args.append(ForeignKey(foreign_key, ondelete=ondelete))
+    kwargs = {
+        "primary_key": primary_key,
+        "nullable": nullable,
+        "index": index,
+        "unique": unique,
+    }
+
+    sa_default = Undefined
+    if field_info.default_factory:
+        sa_default = field_info.default_factory
+    elif field_info.default is not Undefined:
+        sa_default = field_info.default
+    if sa_default is not Undefined:
+        kwargs["default"] = sa_default
+
+    sa_column_args = getattr(field_info, "sa_column_args", Undefined)
+    if sa_column_args is not Undefined:
+        args.extend(list(cast(Sequence[Any], sa_column_args)))
+
+    sa_column_kwargs = getattr(field_info, "sa_column_kwargs", Undefined)
+
+    # IMPORTANT: change from the original function
+    if field_info.description:
+        if sa_column_kwargs is Undefined:
+            sa_column_kwargs = {}
+
+        assert isinstance(sa_column_kwargs, dict)
+
+        # only update comments if not already set
+        if "comment" not in sa_column_kwargs:
+            sa_column_kwargs["comment"] = field_info.description
+
+    if sa_column_kwargs is not Undefined:
+        kwargs.update(cast(Dict[Any, Any], sa_column_kwargs))
+
+    sa_type = get_sqlalchemy_type(field)
+    return Column(sa_type, *args, **kwargs)  # type: ignore
+
+
+sqlmodel.main.get_column_from_field = get_column_from_field

activemodel/logger.py

@@ -0,0 +1,3 @@
+import logging
+
+logger = logging.getLogger(__name__)

activemodel/mixins/__init__.py

@@ -0,0 +1,4 @@
+from .pydantic_json import PydanticJSONMixin
+from .soft_delete import SoftDeletionMixin
+from .timestamps import TimestampsMixin
+from .typeid import TypeIDMixin

activemodel/mixins/pydantic_json.py

@@ -0,0 +1,69 @@
+from types import UnionType
+from typing import get_args, get_origin
+
+from pydantic import BaseModel as PydanticBaseModel
+from sqlalchemy.orm import reconstructor
+
+
+class PydanticJSONMixin:
+    """
+    By default, SQLModel does not convert JSONB columns into pydantic models when they are loaded from the database.
+
+    This mixin, combined with a custom serializer, fixes that issue.
+    """
+
+    @reconstructor
+    def init_on_load(self):
+        # TODO do we need to inspect sa_type
+        for field_name, field_info in self.model_fields.items():
+            raw_value = getattr(self, field_name, None)
+
+            if raw_value is None:
+                continue
+
+            annotation = field_info.annotation
+            origin = get_origin(annotation)
+
+            # e.g. `dict` or `dict[str, str]`, we don't want to do anything with these
+            if origin is dict:
+                continue
+
+            annotation_args = get_args(annotation)
+            is_top_level_list = origin is list
+
+            # if origin is not None:
+            #     assert annotation.__class__ == origin
+
+            model_cls = annotation
+
+            # e.g. SomePydanticModel | None or list[SomePydanticModel] | None
+            # annotation_args are (type, NoneType) in this case
+            if isinstance(annotation, UnionType):
+                non_none_types = [t for t in annotation_args if t is not type(None)]
+
+                if len(non_none_types) == 1:
+                    model_cls = non_none_types[0]
+
+            # e.g. list[SomePydanticModel] | None, we have to unpack it
+            # model_cls will print as a list, but it contains a subtype if you dig into it
+            if (
+                get_origin(model_cls) is list
+                and len(list_annotation_args := get_args(model_cls)) == 1
+            ):
+                model_cls = list_annotation_args[0]
+                is_top_level_list = True
+
+            # e.g. list[SomePydanticModel] or list[SomePydanticModel] | None
+            # iterate through the list and run each item through the pydantic model
+            if is_top_level_list:
+                if isinstance(raw_value, list) and issubclass(
+                    model_cls, PydanticBaseModel
+                ):
+                    parsed_value = [model_cls(**item) for item in raw_value]
+                    setattr(self, field_name, parsed_value)
+
+                continue
+
+            # single class
+            if issubclass(model_cls, PydanticBaseModel):
+                setattr(self, field_name, model_cls(**raw_value))

activemodel/mixins/soft_delete.py

@@ -0,0 +1,17 @@
+from datetime import datetime
+
+import sqlalchemy as sa
+from sqlmodel import Field
+
+
+class SoftDeletionMixin:
+    deleted_at: datetime = Field(
+        default=None,
+        nullable=True,
+        # TODO https://github.com/fastapi/sqlmodel/discussions/1228
+        sa_type=sa.DateTime(timezone=True),  # type: ignore
+    )
+
+    def soft_delete(self):
+        self.deleted_at = datetime.now()
+        raise NotImplementedError("Soft deletion is not implemented")

activemodel/{timestamps.py → mixins/timestamps.py}

@@ -1,10 +1,9 @@
 from datetime import datetime
 
 import sqlalchemy as sa
-
-# TODO raw sql https://github.com/tiangolo/sqlmodel/discussions/772
 from sqlmodel import Field
 
+# TODO raw sql https://github.com/tiangolo/sqlmodel/discussions/772
 # @classmethod
 # def select(cls):
 #     with get_session() as session:
@@ -14,11 +13,11 @@ from sqlmodel import Field
 #             yield result
 
 
-class TimestampMixin:
+class TimestampsMixin:
     """
     Simple created at and updated at timestamps. Mix them into your model:
 
-    >>> class MyModel(TimestampMixin, SQLModel):
+    >>> class MyModel(TimestampsMixin, SQLModel):
     >>>    pass
 
     Originally pulled from: https://github.com/tiangolo/sqlmodel/issues/252