-
Notifications
You must be signed in to change notification settings - Fork 6
/
manager.py
114 lines (89 loc) · 3.57 KB
/
manager.py
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
from sqlalchemy import engine_from_config, orm
from sqlalchemy.orm.exc import UnmappedError
import query, model
class ManagerBase(object):
'''Useful extensions to self.session'''
def add(self, *instances):
'''Override `session.add()` so it can function like `session.add_all()` and support chaining'''
for instance in instances:
if isinstance(instance, list):
self.add(*instance)
else:
self.session.add(instance)
return self.session
def add_commit(self, *instances):
'''Add instances to session and commit in one call'''
self.add(*instances).commit()
def delete(self, *instances):
'''Override `session.delete()` so it can function like `session.add_all()` and support chaining'''
for instance in instances:
if isinstance(instance, list):
self.delete(*instance)
else:
self.session.delete(instance)
return self.session
def delete_commit(self, *instances):
'''Delete instances to session and commit in one call'''
self.delete(*instances).commit()
class Manager(ManagerBase):
'''
Manager for session
'''
def __init__(self, Model=None, config=None, session=None, engine=None):
self.session = None
# declarative base class
self.Model = Model or self.model.make_declarative_base()
if config is None:
config = {}
self.init_session(config=config.get('session'), session=session)
self.init_engine(config=config.get('engine'), engine=engine)
def init_engine(self, config=None, engine=None, config_prefix=''):
'''
Initialize engine. Allow for lazy configuration after `__init__()`.
If both `config` and `engine` are supplied, then the config generated engine will take precedence.
New engine creation will also result in a new `session` object using `engine`.
'''
if config:
engine = engine_from_config(config, prefix=config_prefix)
self.init_session(config={'bind': engine})
def init_session(self, config=None, session=None):
'''
Initialize session. Allow for lazy configuration after `__init__()`.
If both `config` and `session` are supplied, `config` will configure `session`.
'''
self.session = session
if self.session is None:
self.session = orm.scoped_session(orm.sessionmaker())
if config is None:
config = {}
config.setdefault('autocommit', False)
config.setdefault('autoflush', True)
config.setdefault('query_cls', query.Query)
self.session.configure(**config)
if self.Model:
model.extend_declarative_base(self.Model, self.session)
@property
def metadata(self):
return getattr(self.Model, 'metadata', None)
@property
def engine(self):
return self.session.get_bind()
def create_all(self):
'''
Creates database schema from models
'''
if self.metadata is None:
raise UnmappedError('Missing declarative base model')
self.metadata.create_all(bind=self.engine)
def drop_all(self):
'''
Drops tables defined by models
'''
if self.metadata is None:
raise UnmappedError('Missing declarative base model')
self.metadata.drop_all(bind=self.engine)
def __getattr__(self, attr):
'''
Delegate all other attributes to self.session
'''
return getattr(self.session, attr)