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

activemodel/mixins/typeid.py

@@ -0,0 +1,46 @@
+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: TypeIDType = Field(
+            sa_column=Column(TypeIDType(prefix), primary_key=True, nullable=False),
+            default_factory=lambda: TypeID(prefix),
+        )
+
+    _prefixes.append(prefix)
+
+    return _TypeIDMixin
+
+
+# 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
+
+#     Will automatically have an `id` field with prefix "xyz"
+#     """
+
+#     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),
+#         )

activemodel/pytest/__init__.py

@@ -0,0 +1,2 @@
+from .transaction import database_reset_transaction
+from .truncate import database_reset_truncate

activemodel/pytest/transaction.py

@@ -0,0 +1,51 @@
+from activemodel import SessionManager
+
+
+def database_reset_transaction():
+    """
+    Wrap all database interactions for a given test in a nested transaction and roll it back after the test.
+
+    >>> from activemodel.pytest import database_reset_transaction
+    >>> pytest.fixture(scope="function", autouse=True)(database_reset_transaction)
+
+    References:
+
+    - https://stackoverflow.com/questions/62433018/how-to-make-sqlalchemy-transaction-rollback-drop-tables-it-created
+    - https://aalvarez.me/posts/setting-up-a-sqlalchemy-and-pytest-based-test-suite/
+    - https://github.com/nickjj/docker-flask-example/blob/93af9f4fbf185098ffb1d120ee0693abcd77a38b/test/conftest.py#L77
+    - https://github.com/caiola/vinhos.com/blob/c47d0a5d7a4bf290c1b726561d1e8f5d2ac29bc8/backend/test/conftest.py#L46
+    """
+
+    engine = SessionManager.get_instance().get_engine()
+
+    with engine.begin() as connection:
+        transaction = connection.begin_nested()
+
+        SessionManager.get_instance().session_connection = connection
+
+        try:
+            yield
+        finally:
+            transaction.rollback()
+            # TODO is this necessary?
+            connection.close()
+
+
+# TODO unsure if this adds any value beyond the above approach
+# def database_reset_named_truncation():
+#     start_truncation_query = """
+#         BEGIN;
+#         SAVEPOINT test_truncation_savepoint;
+#     """
+
+#     raw_sql_exec(start_truncation_query)
+
+#     yield
+
+#     end_truncation_query = """
+#         ROLLBACK TO SAVEPOINT test_truncation_savepoint;
+#         RELEASE SAVEPOINT test_truncation_savepoint;
+#         ROLLBACK;
+#     """
+
+#     raw_sql_exec(end_truncation_query)

activemodel/pytest/truncate.py

@@ -0,0 +1,46 @@
+from sqlmodel import SQLModel
+
+from ..logger import logger
+from ..session_manager import get_engine
+
+
+def database_reset_truncate():
+    """
+    Transaction is most likely the better way to go, but there are some scenarios where the session override
+    logic does not work properly and you need to truncate tables back to their original state.
+
+    Here's how to do this once at the start of the test:
+
+    >>> from activemodel.pytest import database_reset_truncation
+    >>> def pytest_configure(config):
+    >>> 	database_reset_truncation()
+
+    Or, if you want to use this as a fixture:
+
+    >>> pytest.fixture(scope="function")(database_reset_truncation)
+    >>> def test_the_thing(database_reset_truncation)
+
+    This approach has a couple of problems:
+
+    * You can't run multiple tests in parallel without separate databases
+    * If you have important seed data and want to truncate those tables, the seed data will be lost
+    """
+
+    logger.info("truncating database")
+
+    # TODO get additonal tables to preserve from config
+    exception_tables = ["alembic_version"]
+
+    assert (
+        SQLModel.metadata.sorted_tables
+    ), "No model metadata. Ensure model metadata is imported before running truncate_db"
+
+    with get_engine().connect() as connection:
+        for table in reversed(SQLModel.metadata.sorted_tables):
+            transaction = connection.begin()
+
+            if table.name not in exception_tables:
+                logger.debug(f"truncating table={table.name}")
+                connection.execute(table.delete())
+
+            transaction.commit()

activemodel/query_wrapper.py

@@ -1,28 +1,26 @@
-from typing import Generic, TypeVar
-
+import sqlmodel as sm
 from sqlmodel.sql.expression import SelectOfScalar
 
-WrappedModelType = TypeVar("WrappedModelType")
-
-
-def compile_sql(target: SelectOfScalar):
-    return str(target.compile(get_engine().connect()))
+from .session_manager import get_session
+from .utils import compile_sql
 
 
-class QueryWrapper(Generic[WrappedModelType]):
+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 = 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
 
@@ -31,6 +29,7 @@ class QueryWrapper(Generic[WrappedModelType]):
             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()
 
@@ -40,8 +39,14 @@ class QueryWrapper(Generic[WrappedModelType]):
             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)
 
@@ -51,19 +56,20 @@ class QueryWrapper(Generic[WrappedModelType]):
 
     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
 
@@ -75,7 +81,7 @@ class QueryWrapper(Generic[WrappedModelType]):
 
     def sql(self):
         """
-        Output the raw SQL of the query
+        Output the raw SQL of the query for debugging
         """
 
         return compile_sql(self.target)

activemodel/session_manager.py

@@ -0,0 +1,132 @@
+"""
+Class to make managing sessions with SQL Model easy. Also provides a common entrypoint to make it easy to mutate the
+database environment when testing.
+"""
+
+import contextlib
+import contextvars
+import json
+import typing as t
+
+from decouple import config
+from pydantic import BaseModel
+from sqlalchemy import Connection, Engine
+from sqlmodel import Session, create_engine
+
+
+def _serialize_pydantic_model(model: BaseModel | list[BaseModel] | None) -> str | None:
+    """
+    Pydantic models do not serialize to JSON. You'll get an error such as:
+
+    'TypeError: Object of type TranscriptEntry is not JSON serializable'
+
+    https://github.com/fastapi/sqlmodel/issues/63#issuecomment-2581016387
+
+    This custom serializer is passed to the DB engine to properly serialize pydantic models to
+    JSON for storage in a JSONB column.
+    """
+
+    # TODO I bet this will fail on lists with mixed types
+
+    if isinstance(model, BaseModel):
+        return model.model_dump_json()
+    if isinstance(model, list):
+        # not everything in a list is a pydantic model
+        def dump_if_model(m):
+            if isinstance(m, BaseModel):
+                return m.model_dump()
+            return m
+
+        return json.dumps([dump_if_model(m) for m in model])
+    else:
+        return json.dumps(model)
+
+
+class SessionManager:
+    _instance: t.ClassVar[t.Optional["SessionManager"]] = None
+
+    session_connection: Connection | None
+    "optionally specify a specific session connection to use for all get_session() calls, useful for testing"
+
+    @classmethod
+    def get_instance(cls, database_url: str | None = None) -> "SessionManager":
+        if cls._instance is None:
+            assert database_url is not None, (
+                "Database URL required for first initialization"
+            )
+            cls._instance = cls(database_url)
+
+        return cls._instance
+
+    def __init__(self, database_url: str):
+        self._database_url = database_url
+        self._engine = None
+
+        self.session_connection = None
+
+    # TODO why is this type not reimported?
+    def get_engine(self) -> Engine:
+        if not self._engine:
+            self._engine = create_engine(
+                self._database_url,
+                json_serializer=_serialize_pydantic_model,
+                echo=config("ACTIVEMODEL_LOG_SQL", cast=bool, default=False),
+                # https://docs.sqlalchemy.org/en/20/core/pooling.html#disconnect-handling-pessimistic
+                pool_pre_ping=True,
+                # some implementations include `future=True` but it's not required anymore
+            )
+
+        return self._engine
+
+    def get_session(self):
+        if gsession := _session_context.get():
+
+            @contextlib.contextmanager
+            def _reuse_session():
+                yield gsession
+
+            return _reuse_session()
+
+        if self.session_connection:
+            return Session(bind=self.session_connection)
+
+        return Session(self.get_engine())
+
+
+def init(database_url: str):
+    return SessionManager.get_instance(database_url)
+
+
+def get_engine():
+    return SessionManager.get_instance().get_engine()
+
+
+def get_session():
+    return SessionManager.get_instance().get_session()
+
+
+# contextvars must be at the top-level of a module! You will not get a warning if you don't do this.
+_session_context = contextvars.ContextVar[Session | None](
+    "session_context", default=None
+)
+
+
+@contextlib.contextmanager
+def global_session():
+    with SessionManager.get_instance().get_session() as s:
+        token = _session_context.set(s)
+
+        try:
+            yield
+        finally:
+            _session_context.reset(token)
+
+
+async def aglobal_session():
+    with SessionManager.get_instance().get_session() as s:
+        token = _session_context.set(s)
+
+        try:
+            yield
+        finally:
+            _session_context.reset(token)

activemodel/types/__init__.py

@@ -0,0 +1 @@
+from .typeid import TypeIDType

activemodel/types/typeid.py

@@ -0,0 +1,191 @@
+"""
+Lifted from: https://github.com/akhundMurad/typeid-python/blob/main/examples/sqlalchemy.py
+"""
+
+from typing import Optional
+from uuid import UUID
+
+from pydantic import (
+    GetJsonSchemaHandler,
+)
+from pydantic_core import CoreSchema, core_schema
+from sqlalchemy import types
+from sqlalchemy.util import generic_repr
+from typeid import TypeID
+
+from activemodel.errors import TypeIDValidationError
+
+
+class TypeIDType(types.TypeDecorator):
+    """
+    A SQLAlchemy TypeDecorator that allows storing TypeIDs in the database.
+    The prefix will not be persisted, instead the database-native UUID field will be used.
+    At retrieval time a TypeID will be constructed based on the configured prefix and the
+    UUID value from the database.
+
+    Usage:
+        # will result in TypeIDs such as "user_01h45ytscbebyvny4gc8cr8ma2"
+        id = mapped_column(
+            TypeIDType("user"),
+            primary_key=True,
+            default=lambda: TypeID("user")
+        )
+    """
+
+    impl = types.Uuid
+    # impl = uuid.UUID
+    cache_ok = True
+    prefix: Optional[str] = None
+
+    def __init__(self, prefix: Optional[str], *args, **kwargs):
+        self.prefix = prefix
+        super().__init__(*args, **kwargs)
+
+    def __repr__(self) -> str:
+        # Customize __repr__ to ensure that auto-generated code e.g. from alembic includes
+        # the right __init__ params (otherwise by default prefix will be omitted because
+        # uuid.__init__ does not have such an argument).
+        # TODO this makes it so inspected code does NOT include the suffix
+        return generic_repr(
+            self,
+            to_inspect=TypeID(self.prefix),
+        )
+
+    def process_bind_param(self, value, dialect):
+        """
+        This is run when a search query is built or ...
+        """
+
+        if value is None:
+            return None
+
+        if isinstance(value, UUID):
+            # then it's a UUID class, such as UUID('01942886-7afc-7129-8f57-db09137ed002')
+            return value
+
+        if isinstance(value, str) and value.startswith(self.prefix + "_"):
+            # then it's a TypeID such as 'user_01h45ytscbebyvny4gc8cr8ma2'
+            value = TypeID.from_string(value)
+
+        if isinstance(value, str):
+            # no prefix, raw UUID, let's coerce it into a UUID which SQLAlchemy can handle
+            # ex: '01942886-7afc-7129-8f57-db09137ed002'
+            return UUID(value)
+
+        if isinstance(value, TypeID):
+            # TODO in what case could this None prefix ever occur?
+            if self.prefix is None:
+                if value.prefix is None:
+                    raise TypeIDValidationError(
+                        "Must have a valid prefix set on the class"
+                    )
+            else:
+                if value.prefix != self.prefix:
+                    raise TypeIDValidationError(
+                        f"Expected '{self.prefix}' but got '{value.prefix}'"
+                    )
+
+            return value.uuid
+
+        raise ValueError("Unexpected input type")
+
+    def process_result_value(self, value, dialect):
+        if value is None:
+            return None
+
+        return TypeID.from_uuid(value, self.prefix)
+
+    # def coerce_compared_value(self, op, value):
+    #     """
+    #     This method is called when SQLAlchemy needs to compare a column to a value.
+    #     By returning self, we indicate that this type can handle TypeID instances.
+    #     """
+    #     if isinstance(value, TypeID):
+    #         return self
+
+    #     return super().coerce_compared_value(op, value)
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> CoreSchema:
+        """
+        This fixes the following error: 'Unable to serialize unknown type' by telling pydantic how to serialize this field.
+
+        Note that TypeIDType MUST be the type of the field in SQLModel otherwise you'll get serialization errors.
+        This is done automatically for the mixin but for any relationship fields you'll need to specify the type explicitly.
+
+        - https://github.com/karma-dev-team/karma-system/blob/ee0c1a06ab2cb7aaca6dc4818312e68c5c623365/app/server/value_objects/steam_id.py#L88
+        - https://github.com/hhimanshu/uv-workspaces/blob/main/packages/api/src/_lib/dto/typeid_field.py
+        - https://github.com/karma-dev-team/karma-system/blob/ee0c1a06ab2cb7aaca6dc4818312e68c5c623365/app/base/typeid/type_id.py#L14
+        - https://github.com/pydantic/pydantic/issues/10060
+        - https://github.com/fastapi/fastapi/discussions/10027
+        - https://github.com/alice-biometrics/petisco/blob/b01ef1b84949d156f73919e126ed77aa8e0b48dd/petisco/base/domain/model/uuid.py#L50
+        """
+
+        from_uuid_schema = core_schema.chain_schema(
+            [
+                # TODO not sure how this is different from the UUID schema, should try it  out.
+                # core_schema.is_instance_schema(TypeID),
+                # core_schema.uuid_schema(),
+                core_schema.no_info_plain_validator_function(
+                    TypeID.from_string,
+                    json_schema_input_schema=core_schema.str_schema(),
+                ),
+            ]
+        )
+
+        return core_schema.json_or_python_schema(
+            json_schema=from_uuid_schema,
+            # metadata=core_schema.str_schema(
+            #     pattern="^[0-9a-f]{24}$",
+            #     min_length=24,
+            #     max_length=24,
+            # ),
+            # metadata={
+            #     "pydantic_js_input_core_schema": core_schema.str_schema(
+            #         pattern="^[0-9a-f]{24}$",
+            #         min_length=24,
+            #         max_length=24,
+            #     )
+            # },
+            python_schema=core_schema.union_schema([from_uuid_schema]),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                lambda x: str(x)
+            ),
+        )
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls, schema: CoreSchema, handler: GetJsonSchemaHandler
+    ):
+        """
+        Called when generating the openapi schema. This overrides the `function-plain` type which
+        is generated by the `no_info_plain_validator_function`.
+
+        This logis seems to be a hot part of the codebase, so I'd expect this to break as pydantic
+        fastapi continue to evolve.
+
+        Note that this method can return multiple types. A return value can be as simple as:
+
+        {"type": "string"}
+
+        Or, you could return a more specific JSON schema type:
+
+        core_schema.uuid_schema()
+
+        The problem with using something like uuid_schema is the specifi patterns
+
+        https://github.com/BeanieODM/beanie/blob/2190cd9d1fc047af477d5e6897cc283799f54064/beanie/odm/fields.py#L153
+        """
+
+        return {
+            "type": "string",
+            # TODO implement a more strict pattern in regex
+            #      https://github.com/jetify-com/typeid/blob/3d182feed5687c21bb5ab93d5f457ff96749b68b/spec/README.md?plain=1#L38
+            # "pattern": "^[0-9a-f]{24}$",
+            # "minLength": 24,
+            # "maxLength": 24,
+        }
+
+        return core_schema.uuid_schema()

activemodel/utils.py

@@ -0,0 +1,65 @@
+import inspect
+import pkgutil
+import sys
+from types import ModuleType
+
+from sqlalchemy import text
+from sqlmodel import SQLModel
+from sqlmodel.sql.expression import SelectOfScalar
+
+from .logger import logger
+from .session_manager import get_engine, get_session
+
+
+def compile_sql(target: SelectOfScalar):
+    "convert a query into SQL, helpful for debugging"
+    dialect = get_engine().dialect
+    # TODO I wonder if we could store the dialect to avoid getting an engine reference
+    compiled = target.compile(dialect=dialect, compile_kwargs={"literal_binds": True})
+    return str(compiled)
+
+
+# TODO document further, lots of risks here
+def raw_sql_exec(raw_query: str):
+    with get_session() as session:
+        session.execute(text(raw_query))
+
+
+def find_all_sqlmodels(module: ModuleType):
+    """Import all model classes from module and submodules into current namespace."""
+
+    logger.debug(f"Starting model import from module: {module.__name__}")
+    model_classes = {}
+
+    # Walk through all submodules
+    for loader, module_name, is_pkg in pkgutil.walk_packages(module.__path__):
+        full_name = f"{module.__name__}.{module_name}"
+        logger.debug(f"Importing submodule: {full_name}")
+
+        # Check if module is already imported
+        if full_name in sys.modules:
+            submodule = sys.modules[full_name]
+        else:
+            logger.warning(
+                f"Module not found in sys.modules, not importing: {full_name}"
+            )
+            continue
+
+        # Get all classes from module
+        for name, obj in inspect.getmembers(submodule):
+            if inspect.isclass(obj) and issubclass(obj, SQLModel) and obj != SQLModel:
+                logger.debug(f"Found model class: {name}")
+                model_classes[name] = obj
+
+    logger.debug(f"Completed model import. Found {len(model_classes)} models")
+    return model_classes
+
+
+def hash_function_code(func):
+    "get sha of a function to easily assert that it hasn't changed"
+
+    import hashlib
+    import inspect
+
+    source = inspect.getsource(func)
+    return hashlib.sha256(source.encode()).hexdigest()