activemodel 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

activemodel/pytest/transaction.py

@@ -1,19 +1,107 @@
+import contextlib
+import contextvars
+
+from sqlmodel import Session
 from activemodel import SessionManager
+from activemodel.session_manager import global_session
 
 from ..logger import logger
 
+try:
+    import factory as factory_exists
+except ImportError:
+    factory_exists = None
+
+try:
+    import polyfactory as polyfactory_exists
+except ImportError:
+    polyfactory_exists = None
+
+
+_test_session = contextvars.ContextVar[Session | None]("test_session", default=None)
+
+
+def set_factory_session(session):
+    if not factory_exists:
+        return
+    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"
+
+
+def set_polyfactory_session(session):
+    if not polyfactory_exists:
+        return
+
+    from .factories import ActiveModelFactory
+
+    ActiveModelFactory.__sqlalchemy_session__ = session
+
+
+@contextlib.contextmanager
+def test_session():
+    """
+    Configures a session-global database session for a test.
+
+    Use this as a fixture using `db_session`. This method is meant to be used as a context manager.
+
+    This is useful for tests that need to interact with the database multiple times before calling application code
+    that uses the objects. This is intended to be used outside of an integration test. Integration tests generally
+    do not use database transactions to clean the database and instead use truncation. The transaction fixture
+    configures a session, which is then used here. This method requires that this global test session is already
+    configured. If the transaction fixture is not used, then there is no session available for use and this will fail.
+
+    ActiveModelFactory.save() does this automatically, but if you need to manually create objects
+    and persist them to a DB, you can run into issues with the simple `expunge()` call
+    used to reassociate an object with a new session. If there are more complex relationships
+    this approach will fail and give you detached object errors.
+
+    >>> from activemodel.pytest import test_session
+    >>> def test_the_thing():
+    >>>     with test_session():
+    ...         obj = MyModel(name="test").save()
+    ...         obj2 = MyModelFactory.save()
+
+    More information: https://grok.com/share/bGVnYWN5_c21dd39f-84a7-44cf-a05b-9b26c8febb0b
+    """
+
+    if model_factory_session := _test_session.get():
+        with global_session(model_factory_session) as session:
+            yield session
+    else:
+        raise ValueError("No test session available")
+
+
+def database_truncate_session():
+    """
+    Provides a database session for testing when using a truncation cleaning strategy.
+
+    When not using a transaction cleaning strategy, no global test session is set
+    """
+    with test_session() as session:
+        yield session
+
 
 def database_reset_transaction():
     """
     Wrap all database interactions for a given test in a nested transaction and roll it back after the test.
 
+    This is provided as a function, not a fixture, since you'll need to determine when a integration test is run. Here's
+    an example of how to build a fixture from this method:
+
     >>> from activemodel.pytest import 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.
+    Transaction-based DB cleaning does *not* work if the DB mutations are happening in a separate process because the
+    same session is not shared across python processes. For this scenario, use the truncate method.
 
-    In this case, you should use the truncate.
+    Note that using `fork` as a multiprocess start method is dangerous. Use spawn. This link has more documentation
+    around this topic:
+
+    https://github.com/iloveitaly/python-starter-template/blob/master/app/configuration/lang.py
 
     References:
 
@@ -28,30 +116,29 @@ def database_reset_transaction():
 
     engine = SessionManager.get_instance().get_engine()
 
-    logger.info("starting global database transaction")
+    logger.debug("starting global 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?
+            raise ValueError("global session already set")
 
+        # NOTE we very intentionally do NOT
         SessionManager.get_instance().session_connection = connection
 
         try:
-            with SessionManager.get_instance().get_session() as factory_session:
-                try:
-                    from factory.alchemy import SQLAlchemyModelFactory
+            with SessionManager.get_instance().get_session() as model_factory_session:
+                # set global database sessions for model factories to avoid lazy loading issues
+                set_factory_session(model_factory_session)
+                set_polyfactory_session(model_factory_session)
 
-                    # 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
+                test_session_token = _test_session.set(model_factory_session)
 
-                yield
+                try:
+                    yield
+                finally:
+                    _test_session.reset(test_session_token)
         finally:
             logger.debug("rolling back transaction")
 

activemodel/pytest/truncate.py

@@ -1,10 +1,110 @@
+import os
+from typing import Iterable
+
+import pytest
 from sqlmodel import SQLModel
 
 from ..logger import logger
 from ..session_manager import get_engine
+from pytest import Config
+import typing as t
+
+T = t.TypeVar("T")
+
+
+def _normalize_to_list_of_strings(str_or_list: list[str] | str) -> list[str]:
+    if isinstance(str_or_list, list):
+        return str_or_list
+
+    raw_list = str_or_list.split(",")
+    return [entry.strip() for entry in raw_list if entry and entry.strip()]
+
+
+def _get_pytest_option(
+    config: Config, key: str, *, cast: t.Callable[[t.Any], T] | None = str
+) -> T | None:
+    if not config:
+        return None
+
+    try:
+        val = config.getoption(key)
+    except ValueError:
+        val = None
+
+    if val is None:
+        val = config.getini(key)
+
+    if val is not None:
+        if cast:
+            return cast(val)
+
+        return val
+
+    return None
+
+
+def _normalize_preserve_tables(raw: Iterable[str]) -> list[str]:
+    """Normalize user supplied table list: strip, dedupe (order not preserved).
+
+    Returns a sorted list (case-insensitive sort while preserving original casing
+    for readability in logs).
+    """
+
+    cleaned = {name.strip() for name in raw if name and name.strip()}
+    # deterministic order: casefold sort
+    return sorted(cleaned, key=lambda s: s.casefold())
+
+
+def _get_excluded_tables(
+    pytest_config: Config | None, preserve_tables: list[str] | None
+) -> list[str]:
+    """Resolve list of tables to exclude (i.e. *preserve* / NOT truncate).
+
+    Precedence (lowest -> highest):
+        1. pytest ini option ``activemodel_preserve_tables`` (if available)
+        2. Environment variable ``ACTIVEMODEL_PRESERVE_TABLES`` (comma separated)
+        3. Function argument ``preserve_tables``
+
+            Behavior:
+                    * If user supplies nothing via any channel, defaults to ["alembic_version"].
+                    * Case-insensitive matching during truncation; returned list is normalized
+                        (deduped, sorted) for deterministic logging.
+                    * Emits a warning only when the ini option is *explicitly* specified but empty after normalization.
+    """
+
+    # 1. pytest ini option (registered as type="linelist" -> typically list[str])
+    ini_tables = (
+        _get_pytest_option(
+            pytest_config,
+            "activemodel_preserve_tables",
+            cast=_normalize_to_list_of_strings,
+        )
+        or []
+    )
+
+    # 2. environment variable
+    env_var = os.getenv("ACTIVEMODEL_PRESERVE_TABLES", "")
+    env_tables = _normalize_to_list_of_strings(env_var)
+
+    # 3. function argument
+    arg_tables = preserve_tables or []
+
+    # Consider customization only if any non-empty source provided values OR the function arg explicitly passed
+    combined_raw = [*ini_tables, *env_tables, *arg_tables]
+
+    # if no user customization, force alembic_version
+    if not combined_raw:
+        return ["alembic_version"]
+
+    normalized = _normalize_preserve_tables(combined_raw)
+    logger.debug(f"excluded tables for truncation: {normalized}")
+
+    return normalized
 
 
-def database_reset_truncate():
+def database_reset_truncate(
+    preserve_tables: list[str] | None = None, pytest_config: Config | None = None
+):
     """
     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.
@@ -28,18 +128,19 @@ def database_reset_truncate():
 
     logger.info("truncating database")
 
-    # TODO get additonal tables to preserve from config
-    exception_tables = ["alembic_version"]
+    # Determine excluded (preserved) tables and build case-insensitive lookup set
+    exception_tables = _get_excluded_tables(pytest_config, preserve_tables)
+    exception_lookup = {t.lower() for t in exception_tables}
 
-    assert (
-        SQLModel.metadata.sorted_tables
-    ), "No model metadata. Ensure model metadata is imported before running truncate_db"
+    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:
+            if table.name.lower() not in exception_lookup:
                 logger.debug(f"truncating table={table.name}")
                 connection.execute(table.delete())
 

activemodel/query_wrapper.py

@@ -1,11 +1,13 @@
 import sqlmodel as sm
 from sqlmodel.sql.expression import SelectOfScalar
 
+from activemodel.types.sqlalchemy_protocol import SQLAlchemyQueryMethods
+
 from .session_manager import get_session
 from .utils import compile_sql
 
 
-class QueryWrapper[T: sm.SQLModel]:
+class QueryWrapper[T: sm.SQLModel](SQLAlchemyQueryMethods[T]):
     """
     Make it easy to run queries off of a model
     """
@@ -46,13 +48,20 @@ class QueryWrapper[T: sm.SQLModel]:
         with get_session() as session:
             return session.scalar(sm.select(sm.func.count()).select_from(self.target))
 
+    def scalar(self):
+        """
+        >>>
+        """
+        with get_session() as session:
+            return session.scalar(self.target)
+
     def exec(self):
         with get_session() as session:
             return session.exec(self.target)
 
     def delete(self):
         with get_session() as session:
-            session.delete(self.target)
+            return session.delete(self.target)
 
     def __getattr__(self, name):
         """
@@ -87,4 +96,5 @@ class QueryWrapper[T: sm.SQLModel]:
         return compile_sql(self.target)
 
     def __repr__(self) -> str:
+        # TODO we should improve structure of this a bit more, maybe wrap in <> or something?
         return f"{self.__class__.__name__}: Current SQL:\n{self.sql()}"

activemodel/session_manager.py

@@ -13,6 +13,8 @@ from pydantic import BaseModel
 from sqlalchemy import Connection, Engine
 from sqlmodel import Session, create_engine
 
+ACTIVEMODEL_LOG_SQL = config("ACTIVEMODEL_LOG_SQL", cast=bool, default=False)
+
 
 def _serialize_pydantic_model(model: BaseModel | list[BaseModel] | None) -> str | None:
     """
@@ -50,18 +52,26 @@ class SessionManager:
     "optionally specify a specific session connection to use for all get_session() calls, useful for testing and migrations"
 
     @classmethod
-    def get_instance(cls, database_url: str | None = None) -> "SessionManager":
+    def get_instance(
+        cls,
+        database_url: str | None = None,
+        *,
+        engine_options: dict[str, t.Any] | 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)
+            cls._instance = cls(database_url, engine_options=engine_options)
 
         return cls._instance
 
-    def __init__(self, database_url: str):
+    def __init__(
+        self, database_url: str, *, engine_options: dict[str, t.Any] | None = None
+    ):
         self._database_url = database_url
         self._engine = None
+        self._engine_options: dict = engine_options or {}
 
         self.session_connection = None
 
@@ -72,11 +82,12 @@ class SessionManager:
                 self._database_url,
                 # NOTE very important! This enables pydantic models to be serialized for JSONB columns
                 json_serializer=_serialize_pydantic_model,
-                # TODO move to a constants area
-                echo=config("ACTIVEMODEL_LOG_SQL", cast=bool, default=False),
+                echo=ACTIVEMODEL_LOG_SQL,
+                echo_pool=ACTIVEMODEL_LOG_SQL,
                 # 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
+                **self._engine_options,
             )
 
         return self._engine
@@ -99,9 +110,10 @@ class SessionManager:
         return Session(self.get_engine())
 
 
-def init(database_url: str):
+# TODO would be great one day to type engine_options as the SQLAlchemy EngineOptions
+def init(database_url: str, *, engine_options: dict[str, t.Any] | None = None):
     "configure activemodel to connect to a specific database"
-    return SessionManager.get_instance(database_url)
+    return SessionManager.get_instance(database_url, engine_options=engine_options)
 
 
 def get_engine():
@@ -128,18 +140,40 @@ a place to persist a session to use globally across the application.
 
 
 @contextlib.contextmanager
-def global_session():
+def global_session(session: Session | None = None):
     """
-    Generate a session shared across all activemodel calls.
+    Generate a session and share it across all activemodel calls.
 
     Alternatively, you can pass a session to use globally into the context manager, which is helpful for migrations
     and testing.
+
+    This may only be called a single time per callstack. There is one exception: if you call this multiple times
+    and pass in the same session reference, it will result in a noop.
+
+    Args:
+        session: Use an existing session instead of creating a new one
     """
 
+    if session is not None and _session_context.get() is session:
+        yield session
+        return
+
     if _session_context.get() is not None:
-        raise RuntimeError("global session already set")
+        raise RuntimeError("ActiveModel: global session already set")
 
-    with SessionManager.get_instance().get_session() as s:
+    @contextlib.contextmanager
+    def manage_existing_session():
+        "if an existing session already exists, use it without triggering another __enter__"
+        yield session
+
+    # Use provided session or create a new one
+    session_context = (
+        manage_existing_session()
+        if session is not None
+        else SessionManager.get_instance().get_session()
+    )
+
+    with session_context as s:
         token = _session_context.set(s)
 
         try:

activemodel/types/sqlalchemy_protocol.py

@@ -0,0 +1,10 @@
+# IMPORTANT: This file is auto-generated. Do not edit directly.
+
+from typing import Protocol, TypeVar, Any, Generic
+import sqlmodel as sm
+from sqlalchemy.sql.base import _NoArg
+from typing import TYPE_CHECKING
+
+
+class SQLAlchemyQueryMethods[T: sm.SQLModel](Protocol):
+    pass

activemodel/types/sqlalchemy_protocol.pyi

@@ -0,0 +1,132 @@
+# IMPORTANT: This file is auto-generated. Do not edit directly.
+
+from typing import Protocol, TypeVar, Any, Generic
+import sqlmodel as sm
+from sqlalchemy.sql.base import _NoArg
+
+from ..query_wrapper import QueryWrapper
+
+class SQLAlchemyQueryMethods[T: sm.SQLModel](Protocol):
+    """Protocol defining SQLAlchemy query methods forwarded by QueryWrapper.__getattr__"""
+
+    def add_columns(self, *entities: Any) -> "QueryWrapper[T]": ...
+    def add_cte(self, *ctes: Any, nest_here: Any = False) -> "QueryWrapper[T]": ...
+    def alias(self, name: Any = None, flat: Any = False) -> "QueryWrapper[T]": ...
+    def argument_for(self, argument_name: Any, default: Any) -> "QueryWrapper[T]": ...
+    def as_scalar(
+        self,
+    ) -> "QueryWrapper[T]": ...
+    def column(self, column: Any) -> "QueryWrapper[T]": ...
+    def compare(self, other: Any, **kw: Any) -> "QueryWrapper[T]": ...
+    def compile(
+        self, bind: Any = None, dialect: Any = None, **kw: Any
+    ) -> "QueryWrapper[T]": ...
+    def correlate(self, *fromclauses: Any) -> "QueryWrapper[T]": ...
+    def correlate_except(self, *fromclauses: Any) -> "QueryWrapper[T]": ...
+    def corresponding_column(
+        self, column: Any, require_embedded: Any = False
+    ) -> "QueryWrapper[T]": ...
+    def cte(
+        self, name: Any = None, recursive: Any = False, nesting: Any = False
+    ) -> "QueryWrapper[T]": ...
+    def distinct(self, *expr: Any) -> "QueryWrapper[T]": ...
+    def except_(self, *other: Any) -> "QueryWrapper[T]": ...
+    def except_all(self, *other: Any) -> "QueryWrapper[T]": ...
+    def execution_options(self, **kw: Any) -> "QueryWrapper[T]": ...
+    def exists(
+        self,
+    ) -> "QueryWrapper[T]": ...
+    def fetch(
+        self,
+        count: Any,
+        with_ties: Any = False,
+        percent: Any = False,
+        **dialect_kw: Any,
+    ) -> "QueryWrapper[T]": ...
+    def filter(self, *criteria: Any) -> "QueryWrapper[T]": ...
+    def filter_by(self, **kwargs: Any) -> "QueryWrapper[T]": ...
+    def from_statement(self, statement: Any) -> "QueryWrapper[T]": ...
+    def get_children(self, **kw: Any) -> "QueryWrapper[T]": ...
+    def get_execution_options(
+        self,
+    ) -> "QueryWrapper[T]": ...
+    def get_final_froms(
+        self,
+    ) -> "QueryWrapper[T]": ...
+    def get_label_style(
+        self,
+    ) -> "QueryWrapper[T]": ...
+    def group_by(
+        self, _GenerativeSelect__first: Any = _NoArg.NO_ARG, *clauses: Any
+    ) -> "QueryWrapper[T]": ...
+    def having(self, *having: Any) -> "QueryWrapper[T]": ...
+    def intersect(self, *other: Any) -> "QueryWrapper[T]": ...
+    def intersect_all(self, *other: Any) -> "QueryWrapper[T]": ...
+    def is_derived_from(self, fromclause: Any) -> "QueryWrapper[T]": ...
+    def join(
+        self, target: Any, onclause: Any = None, isouter: Any = False, full: Any = False
+    ) -> "QueryWrapper[T]": ...
+    def join_from(
+        self,
+        from_: Any,
+        target: Any,
+        onclause: Any = None,
+        isouter: Any = False,
+        full: Any = False,
+    ) -> "QueryWrapper[T]": ...
+    def label(self, name: Any) -> "QueryWrapper[T]": ...
+    def lateral(self, name: Any = None) -> "QueryWrapper[T]": ...
+    def limit(self, limit: Any) -> "QueryWrapper[T]": ...
+    def memoized_instancemethod(
+        self,
+    ) -> "QueryWrapper[T]": ...
+    def offset(self, offset: Any) -> "QueryWrapper[T]": ...
+    def options(self, *options: Any) -> "QueryWrapper[T]": ...
+    def order_by(
+        self, _GenerativeSelect__first: Any = _NoArg.NO_ARG, *clauses: Any
+    ) -> QueryWrapper[T]: ...
+    def outerjoin(
+        self, target: Any, onclause: Any = None, full: Any = False
+    ) -> "QueryWrapper[T]": ...
+    def outerjoin_from(
+        self, from_: Any, target: Any, onclause: Any = None, full: Any = False
+    ) -> "QueryWrapper[T]": ...
+    def params(
+        self, _ClauseElement__optionaldict: Any = None, **kwargs: Any
+    ) -> "QueryWrapper[T]": ...
+    def prefix_with(self, *prefixes: Any, dialect: Any = "*") -> "QueryWrapper[T]": ...
+    def reduce_columns(self, only_synonyms: Any = True) -> "QueryWrapper[T]": ...
+    def replace_selectable(self, old: Any, alias: Any) -> "QueryWrapper[T]": ...
+    def scalar_subquery(
+        self,
+    ) -> "QueryWrapper[T]": ...
+    def select(self, *arg: Any, **kw: Any) -> "QueryWrapper[T]": ...
+    def select_from(self, *froms: Any) -> "QueryWrapper[T]": ...
+    def self_group(self, against: Any = None) -> "QueryWrapper[T]": ...
+    def set_label_style(self, style: Any) -> "QueryWrapper[T]": ...
+    def slice(self, start: Any, stop: Any) -> "QueryWrapper[T]": ...
+    def subquery(self, name: Any = None) -> "QueryWrapper[T]": ...
+    def suffix_with(self, *suffixes: Any, dialect: Any = "*") -> "QueryWrapper[T]": ...
+    def union(self, *other: Any) -> "QueryWrapper[T]": ...
+    def union_all(self, *other: Any) -> "QueryWrapper[T]": ...
+    def unique_params(
+        self, _ClauseElement__optionaldict: Any = None, **kwargs: Any
+    ) -> "QueryWrapper[T]": ...
+    def where(self, *whereclause: Any) -> "QueryWrapper[T]": ...
+    def with_for_update(
+        self,
+        nowait: Any = False,
+        read: Any = False,
+        of: Any = None,
+        skip_locked: Any = False,
+        key_share: Any = False,
+    ) -> "QueryWrapper[T]": ...
+    def with_hint(
+        self, selectable: Any, text: Any, dialect_name: Any = "*"
+    ) -> "QueryWrapper[T]": ...
+    def with_only_columns(
+        self, *entities: Any, maintain_column_froms: Any = False, **_Select__kw: Any
+    ) -> "QueryWrapper[T]": ...
+    def with_statement_hint(
+        self, text: Any, dialect_name: Any = "*"
+    ) -> "QueryWrapper[T]": ...

activemodel/types/typeid.py

@@ -76,6 +76,7 @@ class TypeIDType(types.TypeDecorator):
         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'
+            # if an invalid uuid is passed, `ValueError('badly formed hexadecimal UUID string')` will be raised
             return UUID(value)
 
         if isinstance(value, TypeID):

activemodel/utils.py

@@ -1,18 +1,12 @@
-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"
+def compile_sql(target: SelectOfScalar) -> str:
+    "convert a query into SQL, helpful for debugging sqlalchemy/sqlmodel queries"
+
     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})
@@ -25,36 +19,6 @@ def raw_sql_exec(raw_query: str):
         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"