-
Notifications
You must be signed in to change notification settings - Fork 2k
/
__init__.py
356 lines (313 loc) · 10.3 KB
/
__init__.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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
# encoding: utf-8
import warnings
import logging
import re
from time import sleep
from os.path import splitext
from sqlalchemy import MetaData, __version__ as sqav, Table
from sqlalchemy.exc import ProgrammingError
from alembic.command import (
upgrade as alembic_upgrade,
downgrade as alembic_downgrade,
current as alembic_current
)
from alembic.config import Config as AlembicConfig
from ckan.model import meta
from ckan.model.meta import (
Session,
engine_is_sqlite,
engine_is_pg,
)
from ckan.model.core import (
State,
)
from ckan.model.system import (
System,
)
from ckan.model.package import (
Package,
PACKAGE_NAME_MIN_LENGTH,
PACKAGE_NAME_MAX_LENGTH,
PACKAGE_VERSION_MAX_LENGTH,
package_table,
)
from ckan.model.tag import (
Tag,
PackageTag,
MAX_TAG_LENGTH,
MIN_TAG_LENGTH,
tag_table,
package_tag_table,
)
from ckan.model.user import (
User,
user_table,
)
from ckan.model.group import (
Member,
Group,
group_table,
member_table,
)
from ckan.model.group_extra import (
GroupExtra,
group_extra_table,
)
from ckan.model.package_extra import (
PackageExtra,
package_extra_table,
)
from ckan.model.resource import (
Resource,
DictProxy,
resource_table,
)
from ckan.model.resource_view import (
ResourceView,
resource_view_table,
)
from ckan.model.tracking import (
tracking_summary_table,
TrackingSummary,
tracking_raw_table
)
from ckan.model.rating import (
Rating,
MIN_RATING,
MAX_RATING,
)
from ckan.model.package_relationship import (
PackageRelationship,
package_relationship_table,
)
from ckan.model.task_status import (
TaskStatus,
task_status_table,
)
from ckan.model.vocabulary import (
Vocabulary,
VOCABULARY_NAME_MAX_LENGTH,
VOCABULARY_NAME_MIN_LENGTH,
)
from ckan.model.activity import (
Activity,
ActivityDetail,
activity_table,
activity_detail_table,
)
from ckan.model.term_translation import (
term_translation_table,
)
from ckan.model.follower import (
UserFollowingUser,
UserFollowingDataset,
UserFollowingGroup,
)
from ckan.model.system_info import (
system_info_table,
SystemInfo,
get_system_info,
set_system_info,
delete_system_info,
)
from ckan.model.domain_object import (
DomainObjectOperation,
DomainObject,
)
from ckan.model.dashboard import (
Dashboard,
)
import ckan.migration
from ckan.common import config
log = logging.getLogger(__name__)
DB_CONNECT_RETRIES = 10
def init_model(engine):
'''Call me before using any of the tables or classes in the model'''
meta.Session.remove()
meta.Session.configure(bind=engine)
meta.create_local_session.configure(bind=engine)
meta.engine = engine
meta.metadata.bind = engine
# sqlalchemy migrate version table
import sqlalchemy.exc
for i in reversed(range(DB_CONNECT_RETRIES)):
try:
Table('alembic_version', meta.metadata, autoload=True)
break
except sqlalchemy.exc.NoSuchTableError:
break
except sqlalchemy.exc.OperationalError as e:
if 'database system is starting up' in repr(e.orig) and i:
sleep(DB_CONNECT_RETRIES - i)
continue
raise
class Repository():
_repo_path, _dot_repo_name = splitext(ckan.migration.__name__)
migrate_repository = _repo_path + ':' + _dot_repo_name[1:]
# note: tables_created value is not sustained between instantiations
# so only useful for tests. The alternative is to use
# are_tables_created().
tables_created_and_initialised = False
def __init__(self, metadata, session):
self.metadata = metadata
self.session = session
self.commit = session.commit
def commit_and_remove(self):
self.session.commit()
self.session.remove()
def init_db(self):
'''Ensures tables, const data and some default config is created.
This method MUST be run before using CKAN for the first time.
Before this method is run, you can either have a clean db or tables
that may have been setup with either upgrade_db or a previous run of
init_db.
'''
warnings.filterwarnings('ignore', 'SAWarning')
self.session.rollback()
self.session.remove()
# sqlite database needs to be recreated each time as the
# memory database is lost.
if self.metadata.bind.engine.url.drivername == 'sqlite':
# this creates the tables, which isn't required inbetween tests
# that have simply called rebuild_db.
self.create_db()
else:
if not self.tables_created_and_initialised:
self.upgrade_db()
self.tables_created_and_initialised = True
log.info('Database initialised')
def clean_db(self):
self.commit_and_remove()
meta.metadata = MetaData(self.metadata.bind)
with warnings.catch_warnings():
warnings.filterwarnings('ignore', '.*(reflection|tsvector).*')
meta.metadata.reflect()
meta.metadata.drop_all()
self.tables_created_and_initialised = False
log.info('Database tables dropped')
def create_db(self):
'''Ensures tables, const data and some default config is created.
i.e. the same as init_db APART from when running tests, when init_db
has shortcuts.
'''
self.metadata.create_all(bind=self.metadata.bind)
log.info('Database tables created')
def rebuild_db(self):
'''Clean and init the db'''
if self.tables_created_and_initialised:
# just delete data, leaving tables - this is faster
self.delete_all()
else:
# delete tables and data
self.clean_db()
self.session.remove()
self.init_db()
self.session.flush()
log.info('Database rebuilt')
def delete_all(self):
'''Delete all data from all tables.'''
self.session.remove()
## use raw connection for performance
connection = self.session.connection()
if sqav.startswith("0.4"):
tables = self.metadata.table_iterator()
else:
tables = reversed(self.metadata.sorted_tables)
for table in tables:
if table.name == 'alembic_version':
continue
connection.execute('delete from "%s"' % table.name)
self.session.commit()
log.info('Database table data deleted')
def reset_alembic_output(self):
self._alembic_output = []
def add_alembic_output(self, *args):
self._alembic_output.append(args)
def take_alembic_output(self, with_reset=True):
output = self._alembic_output
self._alembic_config = []
return output
def setup_migration_version_control(self):
self.reset_alembic_output()
alembic_config = AlembicConfig()
alembic_config.set_main_option(
"script_location", self.migrate_repository
)
alembic_config.set_main_option(
"sqlalchemy.url", str(self.metadata.bind.url)
)
try:
sqlalchemy_migrate_version = self.metadata.bind.execute(
u'select version from migrate_version'
).scalar()
except ProgrammingError:
sqlalchemy_migrate_version = 0
# this value is used for graceful upgrade from
# sqlalchemy-migrate to alembic
alembic_config.set_main_option(
"sqlalchemy_migrate_version", str(sqlalchemy_migrate_version)
)
# This is an interceptor for alembic output. Otherwise,
# everything will be printed to stdout
alembic_config.print_stdout = self.add_alembic_output
self.alembic_config = alembic_config
def current_version(self):
try:
alembic_current(self.alembic_config)
return self.take_alembic_output()[0][0]
except (TypeError, IndexError):
# alembic is not initialized yet
return 'base'
def downgrade_db(self, version='base'):
self.setup_migration_version_control()
alembic_downgrade(self.alembic_config, version)
log.info(u'CKAN database version set to: %s', version)
def upgrade_db(self, version='head'):
'''Upgrade db using sqlalchemy migrations.
@param version: version to upgrade to (if None upgrade to latest)
'''
_assert_engine_msg = (
u'Database migration - only Postgresql engine supported (not %s).'
) % meta.engine.name
assert meta.engine.name in (
u'postgres', u'postgresql'
), _assert_engine_msg
self.setup_migration_version_control()
version_before = self.current_version()
alembic_upgrade(self.alembic_config, version)
version_after = self.current_version()
if version_after != version_before:
log.info(
u'CKAN database version upgraded: %s -> %s',
version_before,
version_after
)
else:
log.info(u'CKAN database version remains as: %s', version_after)
def are_tables_created(self):
meta.metadata = MetaData(self.metadata.bind)
with warnings.catch_warnings():
warnings.filterwarnings('ignore', '.*(reflection|geometry).*')
meta.metadata.reflect()
return bool(meta.metadata.tables)
repo = Repository(meta.metadata, meta.Session)
def is_id(id_string):
'''Tells the client if the string looks like a revision id or not'''
reg_ex = '^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$'
return bool(re.match(reg_ex, id_string))
def parse_db_config(config_key=u'sqlalchemy.url'):
u''' Takes a config key for a database connection url and parses it into
a dictionary. Expects a url like:
'postgres://tester:pass@localhost/ckantest3'
Returns None if the url could not be parsed.
'''
url = config[config_key]
regex = [
u'^\\s*(?P<db_type>\\w*)', u'://', u'(?P<db_user>[^:]*)', u':?',
u'(?P<db_pass>[^@]*)', u'@', u'(?P<db_host>[^/:]*)', u':?',
u'(?P<db_port>[^/]*)', u'/', u'(?P<db_name>[\\w.-]*)'
]
db_details_match = re.match(u''.join(regex), url)
if not db_details_match:
return
return db_details_match.groupdict()