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

activemodel/__init__.py

@@ -1,2 +1,4 @@
 from .base_model import BaseModel
+
+# from .field import Field
 from .session_manager import SessionManager, get_engine, get_session, init

activemodel/base_model.py

@@ -1,33 +1,38 @@
 import json
 import typing as t
-from contextlib import contextmanager
+from uuid import UUID
 
 import pydash
 import sqlalchemy as sa
 import sqlmodel as sm
 from sqlalchemy import Connection, event
 from sqlalchemy.orm import Mapper, declared_attr
-from sqlmodel import Session, SQLModel, select
+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:
 
-# TODO this does not seem to work with the latest 2.9.x pydantic and sqlmodel
-# https://github.com/SE-Sustainability-OSS/ecodev-core/blob/main/ecodev_core/sqlmodel_utils.py
-class SQLModelWithValidation(SQLModel):
-    """
-    Helper class to ease validation in SQLModel classes with table=True
-    """
+- 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
+"""
 
-    @classmethod
-    def create(cls, **kwargs):
-        """
-        Forces validation to take place, even for SQLModel classes with table=True
-        """
-        return cls(**cls.__bases__[0](**kwargs).model_dump())
+SQLModel.metadata.naming_convention = POSTGRES_INDEXES_NAMING_CONVENTION
 
 
 class BaseModel(SQLModel):
@@ -36,10 +41,14 @@ class BaseModel(SQLModel):
 
     https://github.com/woofz/sqlmodel-basecrud/blob/main/sqlmodel_basecrud/basecrud.py
 
-    {before,after} hooks are modeled after Rails.
+    - {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):
@@ -47,6 +56,14 @@ class BaseModel(SQLModel):
 
         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:
@@ -96,11 +113,37 @@ class BaseModel(SQLModel):
         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.
 
@@ -111,10 +154,35 @@ 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):
         with get_session() as session:
             if old_session := Session.object_session(self):
@@ -145,10 +213,16 @@ class BaseModel(SQLModel):
         """
         Returns the number of records in the database.
         """
-        return get_session().exec(sm.select(sm.func.count()).select_from(cls)).one()
+        with get_session() as session:
+            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):
+    def is_new(self) -> bool:
         return not self._sa_instance_state.has_identity
 
     @classmethod
@@ -186,29 +260,31 @@ class BaseModel(SQLModel):
     #      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?
-    @classmethod
+    # TODO can we type the method signature a bit better?
     # def get(cls, *args: sa.BinaryExpression, **kwargs: t.Any):
+    @classmethod
     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 = []
-        elif len(args) == 1 and isinstance(args[0], TypeID):
-            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)
-        return get_session().exec(statement).first()
+
+        with get_session() as session:
+            return session.exec(statement).first()
 
     @classmethod
     def all(cls):
         with get_session() as session:
-            results = session.exec(sa.sql.select(cls))
+            results = session.exec(sm.select(cls))
 
             # TODO do we need this or can we just return results?
             for result in results:
@@ -222,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/mixins/__init__.py

@@ -1,2 +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/mixins/typeid.py

@@ -1,36 +1,46 @@
-import uuid
-
 from sqlmodel import Column, Field
 from typeid import TypeID
 
 from activemodel.types.typeid import TypeIDType
 
+# global list of prefixes to ensure uniqueness
+_prefixes = []
+
 
 def TypeIDMixin(prefix: str):
+    assert prefix
+    assert prefix not in _prefixes, (
+        f"prefix {prefix} already exists, pick a different one"
+    )
+
     class _TypeIDMixin:
-        id: uuid.UUID = Field(
-            sa_column=Column(TypeIDType(prefix), primary_key=True),
+        id: TypeIDType = Field(
+            sa_column=Column(TypeIDType(prefix), primary_key=True, nullable=False),
             default_factory=lambda: TypeID(prefix),
         )
 
+    _prefixes.append(prefix)
+
     return _TypeIDMixin
 
 
-class TypeIDMixin2:
-    """
-    Mixin class that adds a TypeID primary key to models.
+# TODO not sure if I love the idea of a dynamic class for each mixin as used above
+#      may give this approach another shot in the future
+# class TypeIDMixin2:
+#     """
+#     Mixin class that adds a TypeID primary key to models.
 
 
-    >>>    class MyModel(BaseModel, TypeIDMixin, prefix="xyz", table=True):
-    >>>        name: str
+#     >>>    class MyModel(BaseModel, TypeIDMixin, prefix="xyz", table=True):
+#     >>>        name: str
 
-    Will automatically have an `id` field with prefix "xyz"
-    """
+#     Will automatically have an `id` field with prefix "xyz"
+#     """
 
-    def __init_subclass__(cls, *, prefix: str, **kwargs):
-        super().__init_subclass__(**kwargs)
+#     def __init_subclass__(cls, *, prefix: str, **kwargs):
+#         super().__init_subclass__(**kwargs)
 
-        cls.id: uuid.UUID = Field(
-            sa_column=Column(TypeIDType(prefix), primary_key=True),
-            default_factory=lambda: TypeID(prefix),
-        )
+#         cls.id: uuid.UUID = Field(
+#             sa_column=Column(TypeIDType(prefix), primary_key=True),
+#             default_factory=lambda: TypeID(prefix),
+#         )

activemodel/pytest/truncate.py

@@ -40,7 +40,7 @@ def database_reset_truncate():
             transaction = connection.begin()
 
             if table.name not in exception_tables:
-                logger.debug("truncating table=%s", table.name)
+                logger.debug(f"truncating table={table.name}")
                 connection.execute(table.delete())
 
             transaction.commit()

activemodel/query_wrapper.py

@@ -1,20 +1,26 @@
-import sqlmodel
+import sqlmodel as sm
+from sqlmodel.sql.expression import SelectOfScalar
 
+from .session_manager import get_session
+from .utils import compile_sql
 
-class QueryWrapper[T]:
+
+class QueryWrapper[T: sm.SQLModel]:
     """
     Make it easy to run queries off of a model
     """
 
-    def __init__(self, cls, *args) -> None:
+    target: SelectOfScalar[T]
+
+    def __init__(self, cls: T, *args) -> None:
         # TODO add generics here
         # self.target: SelectOfScalar[T] = sql.select(cls)
 
         if args:
             # very naive, let's assume the args are specific select statements
-            self.target = sqlmodel.sql.select(*args).select_from(cls)
+            self.target = sm.select(*args).select_from(cls)
         else:
-            self.target = sql.select(cls)
+            self.target = sm.select(cls)
 
     # TODO the .exec results should be handled in one shot
 
@@ -23,6 +29,7 @@ class QueryWrapper[T]:
             return session.exec(self.target).first()
 
     def one(self):
+        "requires exactly one result in the dataset"
         with get_session() as session:
             return session.exec(self.target).one()
 
@@ -32,8 +39,14 @@ class QueryWrapper[T]:
             for row in result:
                 yield row
 
+    def count(self):
+        """
+        I did some basic tests
+        """
+        with get_session() as session:
+            return session.scalar(sm.select(sm.func.count()).select_from(self.target))
+
     def exec(self):
-        # TODO do we really need a unique session each time?
         with get_session() as session:
             return session.exec(self.target)
 
@@ -43,19 +56,20 @@ class QueryWrapper[T]:
 
     def __getattr__(self, name):
         """
-        This is called to retrieve the function to execute
+        This implements the magic that forwards function calls to sqlalchemy.
         """
 
         # TODO prefer methods defined in this class
+
         if not hasattr(self.target, name):
             return super().__getattribute__(name)
 
-        attr = getattr(self.target, name)
+        sqlalchemy_target = getattr(self.target, name)
 
-        if callable(attr):
+        if callable(sqlalchemy_target):
 
             def wrapper(*args, **kwargs):
-                result = attr(*args, **kwargs)
+                result = sqlalchemy_target(*args, **kwargs)
                 self.target = result
                 return self