Coverage for /var/srv/projects/api.amasfac.comuna18.com/tmp/venv/lib/python3.9/site-packages/factory/alchemy.py: 1%

1# Copyright: See the LICENSE file. 

2 

3from sqlalchemy.exc import IntegrityError 

4from sqlalchemy.orm.exc import NoResultFound 

5 

6from . import base, errors 

7 

8SESSION_PERSISTENCE_COMMIT = 'commit' 

9SESSION_PERSISTENCE_FLUSH = 'flush' 

10VALID_SESSION_PERSISTENCE_TYPES = [ 

11 None, 

12 SESSION_PERSISTENCE_COMMIT, 

13 SESSION_PERSISTENCE_FLUSH, 

14] 

15 

16 

17class SQLAlchemyOptions(base.FactoryOptions): 

18 def _check_sqlalchemy_session_persistence(self, meta, value): 

19 if value not in VALID_SESSION_PERSISTENCE_TYPES: 

20 raise TypeError( 

21 "%s.sqlalchemy_session_persistence must be one of %s, got %r" % 

22 (meta, VALID_SESSION_PERSISTENCE_TYPES, value) 

23 ) 

24 

25 def _build_default_options(self): 

26 return super()._build_default_options() + [ 

27 base.OptionDefault('sqlalchemy_get_or_create', (), inherit=True), 

28 base.OptionDefault('sqlalchemy_session', None, inherit=True), 

29 base.OptionDefault( 

30 'sqlalchemy_session_persistence', 

31 None, 

32 inherit=True, 

33 checker=self._check_sqlalchemy_session_persistence, 

34 ), 

35 ] 

36 

37 

38class SQLAlchemyModelFactory(base.Factory): 

39 """Factory for SQLAlchemy models. """ 

40 

41 _options_class = SQLAlchemyOptions 

42 

43 class Meta: 

44 abstract = True 

45 

46 @classmethod 

47 def _generate(cls, strategy, params): 

48 # Original params are used in _get_or_create if it cannot build an 

49 # object initially due to an IntegrityError being raised 

50 cls._original_params = params 

51 return super()._generate(strategy, params) 

52 

53 @classmethod 

54 def _get_or_create(cls, model_class, session, args, kwargs): 

55 key_fields = {} 

56 for field in cls._meta.sqlalchemy_get_or_create: 

57 if field not in kwargs: 

58 raise errors.FactoryError( 

59 "sqlalchemy_get_or_create - " 

60 "Unable to find initialization value for '%s' in factory %s" % 

61 (field, cls.__name__)) 

62 key_fields[field] = kwargs.pop(field) 

63 

64 obj = session.query(model_class).filter_by( 

65 *args, **key_fields).one_or_none() 

66 

67 if not obj: 

68 try: 

69 obj = cls._save(model_class, session, args, {**key_fields, **kwargs}) 

70 except IntegrityError as e: 

71 session.rollback() 

72 get_or_create_params = { 

73 lookup: value 

74 for lookup, value in cls._original_params.items() 

75 if lookup in cls._meta.sqlalchemy_get_or_create 

76 } 

77 if get_or_create_params: 

78 try: 

79 obj = session.query(model_class).filter_by( 

80 **get_or_create_params).one() 

81 except NoResultFound: 

82 # Original params are not a valid lookup and triggered a create(), 

83 # that resulted in an IntegrityError. 

84 raise e 

85 else: 

86 raise e 

87 

88 return obj 

89 

90 @classmethod 

91 def _create(cls, model_class, *args, **kwargs): 

92 """Create an instance of the model, and save it to the database.""" 

93 session = cls._meta.sqlalchemy_session 

94 

95 if session is None: 

96 raise RuntimeError("No session provided.") 

97 if cls._meta.sqlalchemy_get_or_create: 

98 return cls._get_or_create(model_class, session, args, kwargs) 

99 return cls._save(model_class, session, args, kwargs) 

100 

101 @classmethod 

102 def _save(cls, model_class, session, args, kwargs): 

103 session_persistence = cls._meta.sqlalchemy_session_persistence 

104 

105 obj = model_class(*args, **kwargs) 

106 session.add(obj) 

107 if session_persistence == SESSION_PERSISTENCE_FLUSH: 

108 session.flush() 

109 elif session_persistence == SESSION_PERSISTENCE_COMMIT: 

110 session.commit() 

111 return obj