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 +25 -9
- activemodel/cli/__init__.py +147 -0
- activemodel/mixins/timestamps.py +4 -1
- activemodel/mixins/typeid.py +10 -2
- activemodel/pytest/__init__.py +1 -1
- activemodel/pytest/factories.py +102 -0
- activemodel/pytest/plugin.py +81 -0
- activemodel/pytest/transaction.py +103 -16
- activemodel/pytest/truncate.py +108 -7
- activemodel/query_wrapper.py +12 -2
- activemodel/session_manager.py +64 -14
- activemodel/types/sqlalchemy_protocol.py +10 -0
- activemodel/types/sqlalchemy_protocol.pyi +132 -0
- activemodel/types/typeid.py +36 -22
- activemodel/types/typeid_patch.py +22 -0
- activemodel/utils.py +3 -39
- {activemodel-0.10.0.dist-info → activemodel-0.12.0.dist-info}/METADATA +40 -16
- activemodel-0.12.0.dist-info/RECORD +30 -0
- activemodel-0.12.0.dist-info/entry_points.txt +2 -0
- activemodel-0.10.0.dist-info/RECORD +0 -24
- activemodel-0.10.0.dist-info/entry_points.txt +0 -2
- {activemodel-0.10.0.dist-info → activemodel-0.12.0.dist-info}/WHEEL +0 -0
- {activemodel-0.10.0.dist-info → activemodel-0.12.0.dist-info}/licenses/LICENSE +0 -0
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
|
|
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
|
-
|
|
163
|
-
|
|
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
|
|
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
|
-
) ->
|
|
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()
|
activemodel/mixins/timestamps.py
CHANGED
|
@@ -20,7 +20,10 @@ class TimestampsMixin:
|
|
|
20
20
|
>>> class MyModel(TimestampsMixin, SQLModel):
|
|
21
21
|
>>> pass
|
|
22
22
|
|
|
23
|
-
|
|
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(
|
activemodel/mixins/typeid.py
CHANGED
|
@@ -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
|
-
#
|
|
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)
|
activemodel/pytest/__init__.py
CHANGED
|
@@ -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
|
|
14
|
-
|
|
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
|
-
|
|
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.
|
|
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
|
-
|
|
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
|
|
44
|
-
|
|
45
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|