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

activemodel/pytest/transaction.py

@@ -1,12 +1,19 @@
 from activemodel import SessionManager
 
+from ..logger import logger
+
 
 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)
+    >>> database_reset_transaction = pytest.fixture(scope="function", autouse=True)(database_reset_transaction)
+
+    Transaction-based DB cleaning does *not* work if the DB mutations are happening in a separate process, which should
+    use spawn, because the same session is not shared across processes. Note that using `fork` is dangerous.
+
+    In this case, you should use the truncate.
 
     References:
 
@@ -14,38 +21,43 @@ def database_reset_transaction():
     - 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
+    - https://stackoverflow.com/questions/64095876/multiprocessing-fork-vs-spawn
+
+    Using a named SAVEPOINT does not give us anything extra, so we are not using it.
     """
 
     engine = SessionManager.get_instance().get_engine()
 
+    logger.info("starting database transaction")
+
     with engine.begin() as connection:
         transaction = connection.begin_nested()
 
+        if SessionManager.get_instance().session_connection is not None:
+            logger.warning("session override already exists")
+            # TODO should we throw an exception here?
+
         SessionManager.get_instance().session_connection = connection
 
         try:
-            yield
+            with SessionManager.get_instance().get_session() as factory_session:
+                try:
+                    from factory.alchemy import SQLAlchemyModelFactory
+
+                    # Ensure that all factories use the same session
+                    for factory in SQLAlchemyModelFactory.__subclasses__():
+                        factory._meta.sqlalchemy_session = factory_session
+                        factory._meta.sqlalchemy_session_persistence = "commit"
+                except ImportError:
+                    pass
+
+                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;
-#     """
+            logger.debug("rolling back transaction")
 
-#     raw_sql_exec(start_truncation_query)
-
-#     yield
+            transaction.rollback()
 
-#     end_truncation_query = """
-#         ROLLBACK TO SAVEPOINT test_truncation_savepoint;
-#         RELEASE SAVEPOINT test_truncation_savepoint;
-#         ROLLBACK;
-#     """
+            # TODO is this necessary? unclear
+            connection.close()
 
-#     raw_sql_exec(end_truncation_query)
+            SessionManager.get_instance().session_connection = None

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
 

activemodel/session_manager.py

@@ -3,24 +3,58 @@ Class to make managing sessions with SQL Model easy. Also provides a common entr
 database environment when testing.
 """
 
+import contextlib
+import contextvars
+import json
 import typing as t
 
 from decouple import config
-from sqlalchemy import Engine
+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
+    "singleton instance of SessionManager"
 
-    session_connection: str
+    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"
+            assert database_url is not None, (
+                "Database URL required for first initialization"
+            )
             cls._instance = cls(database_url)
 
         return cls._instance
@@ -28,6 +62,7 @@ class SessionManager:
     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?
@@ -35,6 +70,8 @@ class SessionManager:
         if not self._engine:
             self._engine = create_engine(
                 self._database_url,
+                # NOTE very important! This enables pydantic models to be serialized for JSONB columns
+                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,
@@ -44,6 +81,15 @@ class SessionManager:
         return self._engine
 
     def get_session(self):
+        if gsession := _session_context.get():
+
+            @contextlib.contextmanager
+            def _reuse_session():
+                yield gsession
+
+            return _reuse_session()
+
+        # a connection can generate nested transactions
         if self.session_connection:
             return Session(bind=self.session_connection)
 
@@ -51,6 +97,7 @@ class SessionManager:
 
 
 def init(database_url: str):
+    "configure activemodel to connect to a specific database"
     return SessionManager.get_instance(database_url)
 
 
@@ -60,3 +107,43 @@ def 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.
+# ContextVar is implemented in C, so it's very special and is both thread-safe and asyncio safe. This variable gives us
+# a place to persist a session to use globally across the application.
+_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 s
+        finally:
+            _session_context.reset(token)
+
+
+async def aglobal_session():
+    """
+    Use this as a fastapi dependency to get a session that is shared across the request:
+
+    >>> APIRouter(
+    >>>     prefix="/internal/v1",
+    >>>     dependencies=[
+    >>>         Depends(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

@@ -3,11 +3,18 @@ Lifted from: https://github.com/akhundMurad/typeid-python/blob/main/examples/sql
 """
 
 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):
     """
@@ -45,12 +52,141 @@ class TypeIDType(types.TypeDecorator):
         )
 
     def process_bind_param(self, value, dialect):
-        if self.prefix is None:
-            assert value.prefix is None
-        else:
-            assert value.prefix == self.prefix
+        """
+        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
 
-        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,
+            # TODO in the the future we could add more exact types
+            # 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

@@ -1,15 +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 activemodel import get_engine
+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()