1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
|
# Copyright: See the LICENSE file.
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm.exc import NoResultFound
from . import base, errors
SESSION_PERSISTENCE_COMMIT = 'commit'
SESSION_PERSISTENCE_FLUSH = 'flush'
VALID_SESSION_PERSISTENCE_TYPES = [
None,
SESSION_PERSISTENCE_COMMIT,
SESSION_PERSISTENCE_FLUSH,
]
class SQLAlchemyOptions(base.FactoryOptions):
def _check_sqlalchemy_session_persistence(self, meta, value):
if value not in VALID_SESSION_PERSISTENCE_TYPES:
raise TypeError(
"%s.sqlalchemy_session_persistence must be one of %s, got %r" %
(meta, VALID_SESSION_PERSISTENCE_TYPES, value)
)
@staticmethod
def _check_has_sqlalchemy_session_set(meta, value):
if value is not None and getattr(meta, "sqlalchemy_session", None) is not None:
raise RuntimeError("Provide either a sqlalchemy_session or a sqlalchemy_session_factory, not both")
def _build_default_options(self):
return super()._build_default_options() + [
base.OptionDefault('sqlalchemy_get_or_create', (), inherit=True),
base.OptionDefault('sqlalchemy_session', None, inherit=True),
base.OptionDefault(
'sqlalchemy_session_factory', None, inherit=True, checker=self._check_has_sqlalchemy_session_set
),
base.OptionDefault(
'sqlalchemy_session_persistence',
None,
inherit=True,
checker=self._check_sqlalchemy_session_persistence,
),
]
class SQLAlchemyModelFactory(base.Factory):
"""Factory for SQLAlchemy models. """
_options_class = SQLAlchemyOptions
_original_params = None
class Meta:
abstract = True
@classmethod
def _generate(cls, strategy, params):
# Original params are used in _get_or_create if it cannot build an
# object initially due to an IntegrityError being raised
cls._original_params = params
return super()._generate(strategy, params)
@classmethod
def _get_or_create(cls, model_class, session, args, kwargs):
key_fields = {}
for field in cls._meta.sqlalchemy_get_or_create:
if field not in kwargs:
raise errors.FactoryError(
"sqlalchemy_get_or_create - "
"Unable to find initialization value for '%s' in factory %s" %
(field, cls.__name__))
key_fields[field] = kwargs.pop(field)
obj = session.query(model_class).filter_by(
*args, **key_fields).one_or_none()
if not obj:
try:
obj = cls._save(model_class, session, args, {**key_fields, **kwargs})
except IntegrityError as e:
session.rollback()
if cls._original_params is None:
raise e
get_or_create_params = {
lookup: value
for lookup, value in cls._original_params.items()
if lookup in cls._meta.sqlalchemy_get_or_create
}
if get_or_create_params:
try:
obj = session.query(model_class).filter_by(
**get_or_create_params).one()
except NoResultFound:
# Original params are not a valid lookup and triggered a create(),
# that resulted in an IntegrityError.
raise e
else:
raise e
return obj
@classmethod
def _create(cls, model_class, *args, **kwargs):
"""Create an instance of the model, and save it to the database."""
session_factory = cls._meta.sqlalchemy_session_factory
if session_factory:
cls._meta.sqlalchemy_session = session_factory()
session = cls._meta.sqlalchemy_session
if session is None:
raise RuntimeError("No session provided.")
if cls._meta.sqlalchemy_get_or_create:
return cls._get_or_create(model_class, session, args, kwargs)
return cls._save(model_class, session, args, kwargs)
@classmethod
def _save(cls, model_class, session, args, kwargs):
session_persistence = cls._meta.sqlalchemy_session_persistence
obj = model_class(*args, **kwargs)
session.add(obj)
if session_persistence == SESSION_PERSISTENCE_FLUSH:
session.flush()
elif session_persistence == SESSION_PERSISTENCE_COMMIT:
session.commit()
return obj
|
