/
__init__.py
419 lines (359 loc) · 13 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
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
# encoding: utf-8
from __future__ import annotations
import warnings
import logging
import os
import re
from time import sleep
from typing import Any, Optional
import sqlalchemy as sa
from sqlalchemy import MetaData, Table, inspect
from alembic.command import (
upgrade as alembic_upgrade,
downgrade as alembic_downgrade,
current as alembic_current
)
from alembic.config import Config as AlembicConfig
import ckan.model.meta as meta
from ckan.model.meta import Session, registry
from ckan.exceptions import CkanConfigurationException
from ckan.model.core import (
State,
)
from ckan.model.system import (
System,
)
from ckan.model.package import (
Package,
PackageMember,
PACKAGE_NAME_MIN_LENGTH,
PACKAGE_NAME_MAX_LENGTH,
PACKAGE_VERSION_MAX_LENGTH,
package_table,
package_member_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,
AnonymousUser
)
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.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.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,
)
from ckan.model.api_token import (
ApiToken,
)
import ckan.migration
from ckan.common import config
from sqlalchemy.engine import Engine
from ckan.types import AlchemySession
__all__ = [
"registry", "Session", "State", "System", "Package", "PackageMember",
"PACKAGE_NAME_MIN_LENGTH", "PACKAGE_NAME_MAX_LENGTH",
"PACKAGE_VERSION_MAX_LENGTH", "package_table", "package_member_table",
"Tag", "PackageTag", "MAX_TAG_LENGTH", "MIN_TAG_LENGTH", "tag_table",
"package_tag_table", "User", "user_table", "AnonymousUser", "Member", "Group",
"group_table", "member_table",
"GroupExtra", "group_extra_table", "PackageExtra", "package_extra_table",
"Resource", "DictProxy", "resource_table",
"ResourceView", "resource_view_table",
"PackageRelationship", "package_relationship_table",
"TaskStatus", "task_status_table",
"Vocabulary", "VOCABULARY_NAME_MAX_LENGTH", "VOCABULARY_NAME_MIN_LENGTH",
"term_translation_table", "UserFollowingUser", "UserFollowingDataset",
"UserFollowingGroup", "system_info_table", "SystemInfo",
"get_system_info", "set_system_info", "delete_system_info",
"DomainObjectOperation", "DomainObject", "Dashboard", "ApiToken",
"init_model", "Repository",
"repo", "is_id", "parse_db_config"
]
log = logging.getLogger(__name__)
DB_CONNECT_RETRIES: int = 10
def init_model(engine: Engine) -> None:
'''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
# sqlalchemy migrate version table
import sqlalchemy.exc
for i in reversed(range(DB_CONNECT_RETRIES)):
try:
Table('alembic_version', meta.metadata, autoload_with=engine)
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
else:
break
def ensure_engine() -> Engine:
"""Return initialized SQLAlchemy engine or raise an error.
This function guarantees that engine is initialized and provides a hint
when someone attempts to use the database before model is properly
initialized.
Prefer using this function instead of direct access to engine via
`meta.engine`.
"""
if not meta.engine:
log.error(
"%s:%s must be called before any interaction with the database",
init_model.__module__, init_model.__name__
)
raise CkanConfigurationException("Model is not initialized")
return meta.engine
class Repository():
metadata: MetaData
session: AlchemySession
commit: Any
_alembic_ini: str = os.path.join(
os.path.dirname(ckan.migration.__file__),
u"alembic.ini"
)
_alembic_output: list[tuple[str, ...]]
# 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: bool = False
def __init__(self, metadata: MetaData, session: AlchemySession) -> None:
self.metadata = metadata
self.session = session
self.commit = session.commit
def commit_and_remove(self) -> None:
self.session.commit()
self.session.remove()
def init_db(self) -> None:
'''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.
'''
self.session.rollback()
self.session.remove()
if not self.tables_created_and_initialised:
self.upgrade_db()
self.tables_created_and_initialised = True
log.info('Database initialised')
def clean_db(self) -> None:
self.commit_and_remove()
meta.metadata = MetaData()
engine = ensure_engine()
with warnings.catch_warnings():
warnings.filterwarnings('ignore', '.*(reflection|tsvector).*')
meta.metadata.reflect(engine)
with engine.begin() as conn:
meta.metadata.drop_all(conn)
self.tables_created_and_initialised = False
log.info('Database tables dropped')
def create_db(self) -> None:
'''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.
'''
with ensure_engine().begin() as conn:
self.metadata.create_all(conn)
log.info('Database tables created')
def rebuild_db(self) -> None:
'''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) -> None:
'''Delete all data from all tables.'''
self.session.remove()
## use raw connection for performance
connection: Any = self.session.connection()
inspector = sa.inspect(connection)
tables = reversed(self.metadata.sorted_tables)
for table in tables:
# `alembic_version` contains current migration version of the
# DB. If we drop this information, next attempt to apply migrations
# will fail. Don't worry about `<PLUGIN>_alembic_version` tables
# created by extensions - CKAN metadata does not track them, so
# they'll never appear in this list.
if table.name == 'alembic_version':
continue
# if custom model imported without migrations applied,
# corresponding table can be missing from DB
if not inspector.has_table(table):
continue
connection.execute(sa.delete(table))
self.session.commit()
log.info('Database table data deleted')
def reset_alembic_output(self) -> None:
self._alembic_output = []
def add_alembic_output(self, text: str, *args: str) -> None:
self._alembic_output.append((text, *args))
def take_alembic_output(self,
with_reset: bool=True) -> list[tuple[str, ...]]:
output = self._alembic_output
if with_reset:
self.reset_alembic_output()
return output
def setup_migration_version_control(self) -> None:
self.reset_alembic_output()
alembic_config = AlembicConfig(self._alembic_ini)
alembic_config.set_main_option(
"sqlalchemy.url", config.get("sqlalchemy.url")
)
engine = ensure_engine()
sqlalchemy_migrate_version = 0
db_inspect = inspect(engine)
if db_inspect.has_table("migrate_version"):
with engine.connect() as conn:
sqlalchemy_migrate_version = conn.execute(
sa.text('select version from migrate_version')
).scalar()
# 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) -> Optional[str]:
"""Returns current revision of the migration repository.
Returns None for plugins that has no migrations and "base" for plugins
that has migrations but none of them were applied. If current revision
is the newest one, ` (head)` suffix added to the result
"""
from alembic.util.exc import CommandError
try:
alembic_current(self.alembic_config)
return self.take_alembic_output()[0][0]
except (TypeError, IndexError):
# alembic is not initialized yet
return 'base'
except CommandError:
# trying to get revision of plugin without migrations
return None
def downgrade_db(self, version: str='base') -> None:
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: str='head') -> None:
'''Upgrade db using sqlalchemy migrations.
@param version: version to upgrade to (if None upgrade to latest)
'''
engine = ensure_engine()
if engine.name not in ('postgres', 'postgresql'):
log.error(
'Only Postgresql engine supported (not %s).',
engine.name,
)
raise CkanConfigurationException(engine.name)
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) -> bool:
meta.metadata = MetaData()
with warnings.catch_warnings():
warnings.filterwarnings('ignore', '.*(reflection|geometry).*')
meta.metadata.reflect(meta.engine)
return bool(meta.metadata.tables)
repo = Repository(meta.metadata, meta.Session)
def is_id(id_string: str) -> bool:
'''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: str=u'sqlalchemy.url') -> Optional[dict[str, str]]:
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 None
return db_details_match.groupdict()