activemodel 0.10.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
activemodel/base_model.py CHANGED
@@ -2,10 +2,11 @@ import json
2
2
  import typing as t
3
3
  from uuid import UUID
4
4
 
5
- import pydash
6
5
  import sqlalchemy as sa
7
6
  import sqlmodel as sm
7
+ import textcase
8
8
  from sqlalchemy import Connection, event
9
+ from sqlalchemy.dialects.postgresql import insert as postgres_insert
9
10
  from sqlalchemy.orm import Mapper, declared_attr
10
11
  from sqlalchemy.orm.attributes import flag_modified as sa_flag_modified
11
12
  from sqlalchemy.orm.base import instance_state
@@ -19,7 +20,6 @@ from . import get_column_from_field_patch # noqa: F401
19
20
  from .logger import logger
20
21
  from .query_wrapper import QueryWrapper
21
22
  from .session_manager import get_session
22
- from sqlalchemy.dialects.postgresql import insert as postgres_insert
23
23
 
24
24
  POSTGRES_INDEXES_NAMING_CONVENTION = {
25
25
  "ix": "%(column_0_label)s_idx",
@@ -152,19 +152,19 @@ class BaseModel(SQLModel):
152
152
  @declared_attr
153
153
  def __tablename__(cls) -> str:
154
154
  """
155
- Automatically generates the table name for the model by converting the class name from camel case to snake case.
156
- This is the recommended format for table names:
155
+ Automatically generates the table name for the model by converting the model's class name from camel case to snake case.
156
+ This is the recommended text case style for table names:
157
157
 
158
158
  https://wiki.postgresql.org/wiki/Don%27t_Do_This#Don.27t_use_upper_case_table_or_column_names
159
159
 
160
- By default, the class is lower cased which makes it harder to read.
160
+ By default, the model's class name is lower cased which makes it harder to read.
161
161
 
162
- Many snake_case libraries struggle with snake case for names like LLMCache, which is why we are using a more
163
- complicated implementation from pydash.
162
+ Also, many text case conversion libraries struggle handling words like "LLMCache", this is why we are using
163
+ a more precise library which processes such acronyms: [`textcase`](https://pypi.org/project/textcase/).
164
164
 
165
165
  https://stackoverflow.com/questions/1175208/elegant-python-function-to-convert-camelcase-to-snake-case
166
166
  """
167
- return pydash.strings.snake_case(cls.__name__)
167
+ return textcase.snake(cls.__name__)
168
168
 
169
169
  @classmethod
170
170
  def foreign_key(cls, **kwargs):
@@ -196,12 +196,13 @@ class BaseModel(SQLModel):
196
196
  "convenience method to avoid having to write .select().where() in order to add conditions"
197
197
  return cls.select().where(*args)
198
198
 
199
+ # TODO we should add an instance method for this as well
199
200
  @classmethod
200
201
  def upsert(
201
202
  cls,
202
203
  data: dict[str, t.Any],
203
204
  unique_by: str | list[str],
204
- ) -> None:
205
+ ) -> t.Self:
205
206
  """
206
207
  This method will insert a new record if it doesn't exist, or update the existing record if it does.
207
208
 
@@ -233,6 +234,8 @@ class BaseModel(SQLModel):
233
234
  return result
234
235
 
235
236
  def delete(self):
237
+ "Delete record completely from the database"
238
+
236
239
  with get_session() as session:
237
240
  if old_session := Session.object_session(self):
238
241
  old_session.expunge(self)
@@ -403,6 +406,19 @@ class BaseModel(SQLModel):
403
406
  with get_session() as session:
404
407
  return session.exec(statement).first()
405
408
 
409
+ @classmethod
410
+ def one_or_none(cls, *args: t.Any, **kwargs: t.Any):
411
+ """
412
+ Gets a single record from the database. Pass an PK ID or a kwarg to filter by.
413
+ Returns None if no record is found. Throws an error if more than one record is found.
414
+ """
415
+
416
+ args, kwargs = cls.__process_filter_args__(*args, **kwargs)
417
+ statement = select(cls).filter(*args).filter_by(**kwargs)
418
+
419
+ with get_session() as session:
420
+ return session.exec(statement).one_or_none()
421
+
406
422
  @classmethod
407
423
  def one(cls, *args: t.Any, **kwargs: t.Any):
408
424
  """
@@ -0,0 +1,147 @@
1
+ """
2
+ This module provides utilities for generating Protocol type definitions for SQLAlchemy's
3
+ SelectOfScalar methods, as well as formatting and fixing Python files using ruff.
4
+ """
5
+
6
+ import inspect
7
+ import logging
8
+ import os
9
+ import subprocess
10
+ from pathlib import Path
11
+ from typing import Any # already imported in header of generated file
12
+
13
+ import sqlmodel as sm
14
+ from sqlmodel.sql.expression import SelectOfScalar
15
+
16
+ from test.test_wrapper import QueryWrapper
17
+
18
+ # Set up logging
19
+ logging.basicConfig(level=logging.DEBUG)
20
+ logger = logging.getLogger(__name__)
21
+
22
+ QUERY_WRAPPER_CLASS_NAME = QueryWrapper.__name__
23
+
24
+
25
+ def format_python_file(file_path: str | Path) -> bool:
26
+ """
27
+ Format a Python file using ruff.
28
+
29
+ Args:
30
+ file_path: Path to the Python file to format
31
+
32
+ Returns:
33
+ bool: True if formatting was successful, False otherwise
34
+ """
35
+ try:
36
+ subprocess.run(["ruff", "format", str(file_path)], check=True)
37
+ logger.info(f"Formatted file using ruff at {file_path}")
38
+ return True
39
+ except subprocess.CalledProcessError as e:
40
+ logger.error(f"Error running ruff to format the file: {e}")
41
+ return False
42
+
43
+
44
+ def fix_python_file(file_path: str | Path) -> bool:
45
+ """
46
+ Fix linting issues in a Python file using ruff.
47
+
48
+ Args:
49
+ file_path: Path to the Python file to fix
50
+
51
+ Returns:
52
+ bool: True if fixing was successful, False otherwise
53
+ """
54
+ try:
55
+ subprocess.run(["ruff", "check", str(file_path), "--fix"], check=True)
56
+ logger.info(f"Fixed linting issues using ruff at {file_path}")
57
+ return True
58
+ except subprocess.CalledProcessError as e:
59
+ logger.error(f"Error running ruff to fix the file: {e}")
60
+ return False
61
+
62
+
63
+ def generate_sqlalchemy_protocol():
64
+ """Generate Protocol type definitions for SQLAlchemy SelectOfScalar methods"""
65
+ logger.info("Starting SQLAlchemy protocol generation")
66
+
67
+ header = """
68
+ # IMPORTANT: This file is auto-generated. Do not edit directly.
69
+
70
+ from typing import Protocol, TypeVar, Any, Generic
71
+ import sqlmodel as sm
72
+ from sqlalchemy.sql.base import _NoArg
73
+
74
+ T = TypeVar('T', bound=sm.SQLModel, covariant=True)
75
+
76
+ class SQLAlchemyQueryMethods(Protocol, Generic[T]):
77
+ \"""Protocol defining SQLAlchemy query methods forwarded by QueryWrapper.__getattr__\"""
78
+ """
79
+ # Initialize output list for generated method signatures
80
+ output: list = []
81
+
82
+ try:
83
+ # Get all methods from SelectOfScalar
84
+ methods = inspect.getmembers(SelectOfScalar)
85
+ logger.debug(f"Discovered {len(methods)} methods from SelectOfScalar")
86
+
87
+ for name, method in methods:
88
+ # Skip private/dunder methods
89
+ if name.startswith("_"):
90
+ continue
91
+
92
+ if not inspect.isfunction(method) and not inspect.ismethod(method):
93
+ logger.debug(f"Skipping non-method: {name}")
94
+ continue
95
+
96
+ logger.debug(f"Processing method: {name}")
97
+ try:
98
+ signature = inspect.signature(method)
99
+ params = []
100
+
101
+ # Process parameters, skipping 'self'
102
+ for param_name, param in list(signature.parameters.items())[1:]:
103
+ if param.kind == param.VAR_POSITIONAL:
104
+ params.append(f"*{param_name}: Any")
105
+ elif param.kind == param.VAR_KEYWORD:
106
+ params.append(f"**{param_name}: Any")
107
+ else:
108
+ if param.default is inspect.Parameter.empty:
109
+ params.append(f"{param_name}: Any")
110
+ else:
111
+ default_repr = repr(param.default)
112
+ params.append(f"{param_name}: Any = {default_repr}")
113
+
114
+ params_str = ", ".join(params)
115
+ output.append(
116
+ f' def {name}(self, {params_str}) -> "{QUERY_WRAPPER_CLASS_NAME}[T]": ...'
117
+ )
118
+ except (ValueError, TypeError) as e:
119
+ logger.warning(f"Could not get signature for {name}: {e}")
120
+ # Some methods might not have proper signatures
121
+ output.append(
122
+ f' def {name}(self, *args: Any, **kwargs: Any) -> "{QUERY_WRAPPER_CLASS_NAME}[T]": ...'
123
+ )
124
+
125
+ # Write the output to a file
126
+ protocol_path = (
127
+ Path(__file__).parent.parent / "types" / "sqlalchemy_protocol.py"
128
+ )
129
+
130
+ # Ensure directory exists
131
+ os.makedirs(protocol_path.parent, exist_ok=True)
132
+
133
+ with open(protocol_path, "w") as f:
134
+ f.write(header + "\n".join(output))
135
+
136
+ logger.info(f"Generated SQLAlchemy protocol at {protocol_path}")
137
+
138
+ # Format and fix the generated file with ruff
139
+ format_python_file(protocol_path)
140
+ fix_python_file(protocol_path)
141
+ except Exception as e:
142
+ logger.error(f"Error generating SQLAlchemy protocol: {e}", exc_info=True)
143
+ raise
144
+
145
+
146
+ if __name__ == "__main__":
147
+ generate_sqlalchemy_protocol()
@@ -20,7 +20,10 @@ class TimestampsMixin:
20
20
  >>> class MyModel(TimestampsMixin, SQLModel):
21
21
  >>> pass
22
22
 
23
- Originally pulled from: https://github.com/tiangolo/sqlmodel/issues/252
23
+ Notes:
24
+
25
+ - Originally pulled from: https://github.com/tiangolo/sqlmodel/issues/252
26
+ - Related issue: https://github.com/fastapi/sqlmodel/issues/539
24
27
  """
25
28
 
26
29
  created_at: datetime | None = Field(
@@ -1,6 +1,7 @@
1
1
  from sqlmodel import Column, Field
2
2
  from typeid import TypeID
3
3
 
4
+ from activemodel.types import typeid_patch # noqa: F401
4
5
  from activemodel.types.typeid import TypeIDType
5
6
 
6
7
  # global list of prefixes to ensure uniqueness
@@ -8,11 +9,15 @@ _prefixes: list[str] = []
8
9
 
9
10
 
10
11
  def TypeIDMixin(prefix: str):
12
+ """
13
+ Mixin that adds a TypeID primary key field to a SQLModel. Specify the prefix to use for the TypeID.
14
+ """
15
+
11
16
  # make sure duplicate prefixes are not used!
12
17
  # NOTE this will cause issues on code reloads
13
18
  assert prefix
14
19
  assert prefix not in _prefixes, (
15
- f"prefix {prefix} already exists, pick a different one"
20
+ f"TypeID prefix '{prefix}' already exists, pick a different one"
16
21
  )
17
22
 
18
23
  class _TypeIDMixin:
@@ -23,9 +28,12 @@ def TypeIDMixin(prefix: str):
23
28
  TypeIDType(prefix),
24
29
  primary_key=True,
25
30
  nullable=False,
31
+ # default on the sa_column level ensures that an ID is generated when creating a new record, even when
32
+ # raw SQLAlchemy operations are used instead of activemodel operations
26
33
  default=lambda: TypeID(prefix),
27
34
  ),
28
- # default_factory=lambda: TypeID(prefix),
35
+ # add a database comment to document the prefix, since it's not stored in the DB otherwise
36
+ description=f"TypeID with prefix: {prefix}",
29
37
  )
30
38
 
31
39
  _prefixes.append(prefix)
@@ -1,2 +1,2 @@
1
- from .transaction import database_reset_transaction
1
+ from .transaction import database_reset_transaction, test_session
2
2
  from .truncate import database_reset_truncate
@@ -0,0 +1,102 @@
1
+ """
2
+ Notes on polyfactory:
3
+
4
+ 1. is_supported_type validates that the class can be used to generate a factory
5
+ https://github.com/litestar-org/polyfactory/issues/655#issuecomment-2727450854
6
+ """
7
+
8
+ import typing as t
9
+
10
+ from polyfactory.factories.pydantic_factory import ModelFactory
11
+ from polyfactory.field_meta import FieldMeta
12
+ from typeid import TypeID
13
+
14
+ from activemodel.session_manager import global_session
15
+
16
+ # TODO not currently used
17
+ # def type_id_provider(cls, field_meta):
18
+ # # TODO this doesn't work well with __ args:
19
+ # # https://github.com/litestar-org/polyfactory/pull/666/files
20
+ # return str(TypeID("hi"))
21
+
22
+
23
+ # BaseFactory.add_provider(TypeIDType, type_id_provider)
24
+
25
+
26
+ class SQLModelFactory[T](ModelFactory[T]):
27
+ """
28
+ Base factory for SQLModel models:
29
+
30
+ 1. Ability to ignore all relationship fks
31
+ 2. Option to ignore all pks
32
+ """
33
+
34
+ __is_base_factory__ = True
35
+
36
+ @classmethod
37
+ def should_set_field_value(cls, field_meta: FieldMeta, **kwargs: t.Any) -> bool:
38
+ # TODO what is this checking for?
39
+ has_object_override = hasattr(cls, field_meta.name)
40
+
41
+ # TODO this should be more intelligent, it's goal is to detect all of the relationship field and avoid settings them
42
+ if not has_object_override and (
43
+ field_meta.name == "id" or field_meta.name.endswith("_id")
44
+ ):
45
+ return False
46
+
47
+ return super().should_set_field_value(field_meta, **kwargs)
48
+
49
+
50
+ # TODO we need to think through how to handle relationships and autogenerate them
51
+ class ActiveModelFactory[T](SQLModelFactory[T]):
52
+ __is_base_factory__ = True
53
+ __sqlalchemy_session__ = None
54
+
55
+ # TODO we shouldn't have to type this, but `save()` typing is not working
56
+ @classmethod
57
+ def save(cls, *args, **kwargs) -> T:
58
+ """
59
+ Where this gets tricky, is this can be called multiple times within the same callstack. This can happen when
60
+ a factory uses other factories to create relationships.
61
+
62
+ In a truncation strategy, the __sqlalchemy_session__ is set to None.
63
+ """
64
+ with global_session(cls.__sqlalchemy_session__):
65
+ return cls.build(*args, **kwargs).save()
66
+
67
+ @classmethod
68
+ def foreign_key_typeid(cls):
69
+ """
70
+ Return a random type id for the foreign key on this model.
71
+
72
+ This is helpful for generating TypeIDs for testing 404s, parsing, manually settings, etc
73
+ """
74
+ # TODO right now assumes the model is typeid, maybe we should assert against this?
75
+ primary_key_name = cls.__model__.primary_key_column().name
76
+ return TypeID(
77
+ cls.__model__.model_fields[primary_key_name].sa_column.type.prefix
78
+ )
79
+
80
+ @classmethod
81
+ def should_set_field_value(cls, field_meta: FieldMeta, **kwargs: t.Any) -> bool:
82
+ # do not default deleted at mixin to deleted!
83
+ # TODO should be smarter about detecting if the mixin is in place
84
+ if field_meta.name in ["deleted_at", "updated_at", "created_at"]:
85
+ return False
86
+
87
+ return super().should_set_field_value(field_meta, **kwargs)
88
+
89
+ # @classmethod
90
+ # def build(
91
+ # cls,
92
+ # factory_use_construct: bool | None = None,
93
+ # sqlmodel_save: bool = False,
94
+ # **kwargs: t.Any,
95
+ # ) -> T:
96
+ # result = super().build(factory_use_construct=factory_use_construct, **kwargs)
97
+
98
+ # # TODO allow magic dunder method here
99
+ # if sqlmodel_save:
100
+ # result.save()
101
+
102
+ # return result
@@ -0,0 +1,81 @@
1
+ """Pytest plugin integration for activemodel.
2
+
3
+ Currently provides:
4
+
5
+ * ``db_session`` fixture – quick access to a database session (see ``test_session``)
6
+ * ``activemodel_preserve_tables`` ini option – configure tables to preserve when using
7
+ ``database_reset_truncate`` (comma separated list or multiple lines depending on config style)
8
+
9
+ Configuration examples:
10
+
11
+ pytest.ini::
12
+
13
+ [pytest]
14
+ activemodel_preserve_tables = alembic_version,zip_code,seed_table
15
+
16
+ pyproject.toml::
17
+
18
+ [tool.pytest.ini_options]
19
+ activemodel_preserve_tables = [
20
+ "alembic_version",
21
+ "zip_code",
22
+ "seed_table",
23
+ ]
24
+
25
+ The list always implicitly includes ``alembic_version`` even if not specified.
26
+ """
27
+
28
+ from activemodel.session_manager import global_session
29
+ import pytest
30
+
31
+ from .transaction import set_factory_session, set_polyfactory_session, test_session
32
+
33
+
34
+ def pytest_addoption(
35
+ parser: pytest.Parser,
36
+ ) -> None: # pragma: no cover - executed during collection
37
+ """Register custom ini options.
38
+
39
+ We treat this as a *linelist* so pyproject.toml list syntax works. Comma separated works too because
40
+ pytest splits lines first; users can still provide one line with commas.
41
+ """
42
+
43
+ parser.addini(
44
+ "activemodel_preserve_tables",
45
+ help=(
46
+ "Tables to preserve when calling activemodel.pytest.database_reset_truncate. "
47
+ ),
48
+ type="linelist",
49
+ default=["alembic_version"],
50
+ )
51
+
52
+
53
+ @pytest.fixture(scope="function")
54
+ def db_session():
55
+ """
56
+ Helpful for tests that are more similar to unit tests. If you doing a routing or integration test, you
57
+ probably don't need this. If your unit test is simple (you are just creating a couple of models) you
58
+ can most likely skip this.
59
+
60
+ This is helpful if you are doing a lot of lazy-loaded params or need a database session to be in place
61
+ for testing code that will run within a celery worker or something similar.
62
+
63
+ >>> def the_test(db_session):
64
+ """
65
+ with test_session() as session:
66
+ yield session
67
+
68
+
69
+ @pytest.fixture(scope="function")
70
+ def db_truncate_session():
71
+ """
72
+ Provides a database session for testing when using a truncation cleaning strategy.
73
+
74
+ When not using a transaction cleaning strategy, no global test session is set
75
+ """
76
+ with global_session() as session:
77
+ # set global database sessions for model factories to avoid lazy loading issues
78
+ set_factory_session(session)
79
+ set_polyfactory_session(session)
80
+
81
+ yield session
@@ -1,19 +1,107 @@
1
+ import contextlib
2
+ import contextvars
3
+
4
+ from sqlmodel import Session
1
5
  from activemodel import SessionManager
6
+ from activemodel.session_manager import global_session
2
7
 
3
8
  from ..logger import logger
4
9
 
10
+ try:
11
+ import factory as factory_exists
12
+ except ImportError:
13
+ factory_exists = None
14
+
15
+ try:
16
+ import polyfactory as polyfactory_exists
17
+ except ImportError:
18
+ polyfactory_exists = None
19
+
20
+
21
+ _test_session = contextvars.ContextVar[Session | None]("test_session", default=None)
22
+
23
+
24
+ def set_factory_session(session):
25
+ if not factory_exists:
26
+ return
27
+ from factory.alchemy import SQLAlchemyModelFactory
28
+
29
+ # Ensure that all factories use the same session
30
+ for factory in SQLAlchemyModelFactory.__subclasses__():
31
+ factory._meta.sqlalchemy_session = factory_session
32
+ factory._meta.sqlalchemy_session_persistence = "commit"
33
+
34
+
35
+ def set_polyfactory_session(session):
36
+ if not polyfactory_exists:
37
+ return
38
+
39
+ from .factories import ActiveModelFactory
40
+
41
+ ActiveModelFactory.__sqlalchemy_session__ = session
42
+
43
+
44
+ @contextlib.contextmanager
45
+ def test_session():
46
+ """
47
+ Configures a session-global database session for a test.
48
+
49
+ Use this as a fixture using `db_session`. This method is meant to be used as a context manager.
50
+
51
+ This is useful for tests that need to interact with the database multiple times before calling application code
52
+ that uses the objects. This is intended to be used outside of an integration test. Integration tests generally
53
+ do not use database transactions to clean the database and instead use truncation. The transaction fixture
54
+ configures a session, which is then used here. This method requires that this global test session is already
55
+ configured. If the transaction fixture is not used, then there is no session available for use and this will fail.
56
+
57
+ ActiveModelFactory.save() does this automatically, but if you need to manually create objects
58
+ and persist them to a DB, you can run into issues with the simple `expunge()` call
59
+ used to reassociate an object with a new session. If there are more complex relationships
60
+ this approach will fail and give you detached object errors.
61
+
62
+ >>> from activemodel.pytest import test_session
63
+ >>> def test_the_thing():
64
+ >>> with test_session():
65
+ ... obj = MyModel(name="test").save()
66
+ ... obj2 = MyModelFactory.save()
67
+
68
+ More information: https://grok.com/share/bGVnYWN5_c21dd39f-84a7-44cf-a05b-9b26c8febb0b
69
+ """
70
+
71
+ if model_factory_session := _test_session.get():
72
+ with global_session(model_factory_session) as session:
73
+ yield session
74
+ else:
75
+ raise ValueError("No test session available")
76
+
77
+
78
+ def database_truncate_session():
79
+ """
80
+ Provides a database session for testing when using a truncation cleaning strategy.
81
+
82
+ When not using a transaction cleaning strategy, no global test session is set
83
+ """
84
+ with test_session() as session:
85
+ yield session
86
+
5
87
 
6
88
  def database_reset_transaction():
7
89
  """
8
90
  Wrap all database interactions for a given test in a nested transaction and roll it back after the test.
9
91
 
92
+ This is provided as a function, not a fixture, since you'll need to determine when a integration test is run. Here's
93
+ an example of how to build a fixture from this method:
94
+
10
95
  >>> from activemodel.pytest import database_reset_transaction
11
96
  >>> database_reset_transaction = pytest.fixture(scope="function", autouse=True)(database_reset_transaction)
12
97
 
13
- Transaction-based DB cleaning does *not* work if the DB mutations are happening in a separate process, which should
14
- use spawn, because the same session is not shared across processes. Note that using `fork` is dangerous.
98
+ Transaction-based DB cleaning does *not* work if the DB mutations are happening in a separate process because the
99
+ same session is not shared across python processes. For this scenario, use the truncate method.
15
100
 
16
- In this case, you should use the truncate.
101
+ Note that using `fork` as a multiprocess start method is dangerous. Use spawn. This link has more documentation
102
+ around this topic:
103
+
104
+ https://github.com/iloveitaly/python-starter-template/blob/master/app/configuration/lang.py
17
105
 
18
106
  References:
19
107
 
@@ -28,30 +116,29 @@ def database_reset_transaction():
28
116
 
29
117
  engine = SessionManager.get_instance().get_engine()
30
118
 
31
- logger.info("starting global database transaction")
119
+ logger.debug("starting global database transaction")
32
120
 
33
121
  with engine.begin() as connection:
34
122
  transaction = connection.begin_nested()
35
123
 
36
124
  if SessionManager.get_instance().session_connection is not None:
37
- logger.warning("session override already exists")
38
- # TODO should we throw an exception here?
125
+ raise ValueError("global session already set")
39
126
 
127
+ # NOTE we very intentionally do NOT
40
128
  SessionManager.get_instance().session_connection = connection
41
129
 
42
130
  try:
43
- with SessionManager.get_instance().get_session() as factory_session:
44
- try:
45
- from factory.alchemy import SQLAlchemyModelFactory
131
+ with SessionManager.get_instance().get_session() as model_factory_session:
132
+ # set global database sessions for model factories to avoid lazy loading issues
133
+ set_factory_session(model_factory_session)
134
+ set_polyfactory_session(model_factory_session)
46
135
 
47
- # Ensure that all factories use the same session
48
- for factory in SQLAlchemyModelFactory.__subclasses__():
49
- factory._meta.sqlalchemy_session = factory_session
50
- factory._meta.sqlalchemy_session_persistence = "commit"
51
- except ImportError:
52
- pass
136
+ test_session_token = _test_session.set(model_factory_session)
53
137
 
54
- yield
138
+ try:
139
+ yield
140
+ finally:
141
+ _test_session.reset(test_session_token)
55
142
  finally:
56
143
  logger.debug("rolling back transaction")
57
144