1 """Database management."""
2 from __future__ import annotations
4 from os.path import isfile
5 from difflib import Differ
6 from sqlite3 import connect as sql_connect, Cursor, Row
7 from typing import Any, Self, TypeVar, Generic
8 from plomtask.exceptions import HandledException, NotFoundException
9 from plomtask.dating import valid_date
11 EXPECTED_DB_VERSION = 5
12 MIGRATIONS_DIR = 'migrations'
13 FILENAME_DB_SCHEMA = f'init_{EXPECTED_DB_VERSION}.sql'
14 PATH_DB_SCHEMA = f'{MIGRATIONS_DIR}/{FILENAME_DB_SCHEMA}'
17 class UnmigratedDbException(HandledException):
18 """To identify case of unmigrated DB file."""
22 """Represents the sqlite3 database's file."""
23 # pylint: disable=too-few-public-methods
25 def __init__(self, path: str) -> None:
30 def create_at(cls, path: str) -> DatabaseFile:
31 """Make new DB file at path."""
32 with sql_connect(path) as conn:
33 with open(PATH_DB_SCHEMA, 'r', encoding='utf-8') as f:
34 conn.executescript(f.read())
35 conn.execute(f'PRAGMA user_version = {EXPECTED_DB_VERSION}')
39 def migrate(cls, path: str) -> DatabaseFile:
40 """Apply migrations from_version to EXPECTED_DB_VERSION."""
41 migrations = cls._available_migrations()
42 from_version = cls._get_version_of_db(path)
43 migrations_todo = migrations[from_version+1:]
44 for j, filename in enumerate(migrations_todo):
45 with sql_connect(path) as conn:
46 with open(f'{MIGRATIONS_DIR}/{filename}', 'r',
47 encoding='utf-8') as f:
48 conn.executescript(f.read())
49 user_version = from_version + j + 1
50 with sql_connect(path) as conn:
51 conn.execute(f'PRAGMA user_version = {user_version}')
54 def _check(self) -> None:
55 """Check file exists, and is of proper DB version and schema."""
56 if not isfile(self.path):
57 raise NotFoundException
58 if self._user_version != EXPECTED_DB_VERSION:
59 raise UnmigratedDbException()
60 self._validate_schema()
63 def _available_migrations() -> list[str]:
64 """Validate migrations directory and return sorted entries."""
65 msg_too_big = 'Migration directory points beyond expected DB version.'
66 msg_bad_entry = 'Migration directory contains unexpected entry: '
67 msg_missing = 'Migration directory misses migration of number: '
69 for entry in listdir(MIGRATIONS_DIR):
70 if entry == FILENAME_DB_SCHEMA:
72 toks = entry.split('_', 1)
74 raise HandledException(msg_bad_entry + entry)
77 except ValueError as e:
78 raise HandledException(msg_bad_entry + entry) from e
79 if i > EXPECTED_DB_VERSION:
80 raise HandledException(msg_too_big)
81 migrations[i] = toks[1]
83 for i in range(EXPECTED_DB_VERSION + 1):
84 if i not in migrations:
85 raise HandledException(msg_missing + str(i))
86 migrations_list += [f'{i}_{migrations[i]}']
87 return migrations_list
90 def _get_version_of_db(path: str) -> int:
91 """Get DB user_version, fail if outside expected range."""
92 sql_for_db_version = 'PRAGMA user_version'
93 with sql_connect(path) as conn:
94 db_version = list(conn.execute(sql_for_db_version))[0][0]
95 if db_version > EXPECTED_DB_VERSION:
96 msg = f'Wrong DB version, expected '\
97 f'{EXPECTED_DB_VERSION}, got unknown {db_version}.'
98 raise HandledException(msg)
99 assert isinstance(db_version, int)
103 def _user_version(self) -> int:
104 """Get DB user_version."""
105 return self._get_version_of_db(self.path)
107 def _validate_schema(self) -> None:
108 """Compare found schema with what's stored at PATH_DB_SCHEMA."""
110 def reformat_rows(rows: list[str]) -> list[str]:
114 for subrow in row.split('\n'):
115 subrow = subrow.rstrip()
118 for i, c in enumerate(subrow):
123 elif ',' == c and 0 == in_parentheses:
127 segment = subrow[prev_split:i].strip()
129 new_row += [f' {segment}']
131 segment = subrow[prev_split:].strip()
133 new_row += [f' {segment}']
134 new_row[0] = new_row[0].lstrip()
135 new_row[-1] = new_row[-1].lstrip()
136 if new_row[-1] != ')' and new_row[-3][-1] != ',':
137 new_row[-3] = new_row[-3] + ','
138 new_row[-2:] = [' ' + new_row[-1][:-1]] + [')']
139 new_rows += ['\n'.join(new_row)]
142 sql_for_schema = 'SELECT sql FROM sqlite_master ORDER BY sql'
143 msg_err = 'Database has wrong tables schema. Diff:\n'
144 with sql_connect(self.path) as conn:
145 schema_rows = [r[0] for r in conn.execute(sql_for_schema) if r[0]]
146 schema_rows = reformat_rows(schema_rows)
147 retrieved_schema = ';\n'.join(schema_rows) + ';'
148 with open(PATH_DB_SCHEMA, 'r', encoding='utf-8') as f:
149 stored_schema = f.read().rstrip()
150 if stored_schema != retrieved_schema:
151 diff_msg = Differ().compare(retrieved_schema.splitlines(),
152 stored_schema.splitlines())
153 raise HandledException(msg_err + '\n'.join(diff_msg))
156 class DatabaseConnection:
157 """A single connection to the database."""
159 def __init__(self, db_file: DatabaseFile) -> None:
160 self.conn = sql_connect(db_file.path)
162 def commit(self) -> None:
163 """Commit SQL transaction."""
166 def exec(self, code: str, inputs: tuple[Any, ...] = tuple()) -> Cursor:
167 """Add commands to SQL transaction."""
168 return self.conn.execute(code, inputs)
170 def exec_on_vals(self, code: str, inputs: tuple[Any, ...]) -> Cursor:
171 """Wrapper around .exec appending adequate " (?, …)" to code."""
172 q_marks_from_values = '(' + ','.join(['?'] * len(inputs)) + ')'
173 return self.exec(f'{code} {q_marks_from_values}', inputs)
175 def close(self) -> None:
176 """Close DB connection."""
179 def rewrite_relations(self, table_name: str, key: str, target: int | str,
180 rows: list[list[Any]], key_index: int = 0) -> None:
181 # pylint: disable=too-many-arguments
182 """Rewrite relations in table_name to target, with rows values.
184 Note that single rows are expected without the column and value
185 identified by key and target, which are inserted inside the function
188 self.delete_where(table_name, key, target)
190 values = tuple(row[:key_index] + [target] + row[key_index:])
191 self.exec_on_vals(f'INSERT INTO {table_name} VALUES', values)
193 def row_where(self, table_name: str, key: str,
194 target: int | str) -> list[Row]:
195 """Return list of Rows at table where key == target."""
196 return list(self.exec(f'SELECT * FROM {table_name} WHERE {key} = ?',
199 # def column_where_pattern(self,
203 # keys: list[str]) -> list[Any]:
204 # """Return column of rows where one of keys matches pattern."""
205 # targets = tuple([f'%{pattern}%'] * len(keys))
206 # haystack = ' OR '.join([f'{k} LIKE ?' for k in keys])
207 # sql = f'SELECT {column} FROM {table_name} WHERE {haystack}'
208 # return [row[0] for row in self.exec(sql, targets)]
210 def column_where(self, table_name: str, column: str, key: str,
211 target: int | str) -> list[Any]:
212 """Return column of table where key == target."""
213 return [row[0] for row in
214 self.exec(f'SELECT {column} FROM {table_name} '
215 f'WHERE {key} = ?', (target,))]
217 def column_all(self, table_name: str, column: str) -> list[Any]:
218 """Return complete column of table."""
219 return [row[0] for row in
220 self.exec(f'SELECT {column} FROM {table_name}')]
222 def delete_where(self, table_name: str, key: str,
223 target: int | str) -> None:
224 """Delete from table where key == target."""
225 self.exec(f'DELETE FROM {table_name} WHERE {key} = ?', (target,))
228 BaseModelId = TypeVar('BaseModelId', int, str)
229 BaseModelInstance = TypeVar('BaseModelInstance', bound='BaseModel[Any]')
232 class BaseModel(Generic[BaseModelId]):
233 """Template for most of the models we use/derive from the DB."""
235 to_save: list[str] = []
236 to_save_versioned: list[str] = []
237 to_save_relations: list[tuple[str, str, str, int]] = []
238 id_: None | BaseModelId
239 cache_: dict[BaseModelId, Self]
240 to_search: list[str] = []
243 def __init__(self, id_: BaseModelId | None) -> None:
244 if isinstance(id_, int) and id_ < 1:
245 msg = f'illegal {self.__class__.__name__} ID, must be >=1: {id_}'
246 raise HandledException(msg)
247 if isinstance(id_, str) and "" == id_:
248 msg = f'illegal {self.__class__.__name__} ID, must be non-empty'
249 raise HandledException(msg)
252 def __hash__(self) -> int:
253 hashable = [self.id_] + [getattr(self, name) for name in self.to_save]
254 for definition in self.to_save_relations:
255 attr = getattr(self, definition[2])
256 hashable += [tuple(rel.id_ for rel in attr)]
257 for name in self.to_save_versioned:
258 hashable += [hash(getattr(self, name))]
259 return hash(tuple(hashable))
261 def __eq__(self, other: object) -> bool:
262 if not isinstance(other, self.__class__):
264 return hash(self) == hash(other)
266 def __lt__(self, other: Any) -> bool:
267 if not isinstance(other, self.__class__):
268 msg = 'cannot compare to object of different class'
269 raise HandledException(msg)
270 assert isinstance(self.id_, int)
271 assert isinstance(other.id_, int)
272 return self.id_ < other.id_
275 def as_dict(self) -> dict[str, object]:
276 """Return self as (json.dumps-coompatible) dict."""
277 d: dict[str, object] = {'id': self.id_}
278 for k in self.to_save:
279 attr = getattr(self, k)
280 if hasattr(attr, 'as_dict'):
283 for k in self.to_save_versioned:
284 attr = getattr(self, k)
286 for r in self.to_save_relations:
288 d[attr_name] = [x.as_dict for x in getattr(self, attr_name)]
292 # (we primarily use the cache to ensure we work on the same object in
293 # memory no matter where and how we retrieve it, e.g. we don't want
294 # .by_id() calls to create a new object each time, but rather a pointer
295 # to the one already instantiated)
297 def __getattribute__(self, name: str) -> Any:
298 """Ensure fail if ._disappear() was called, except to check ._exists"""
299 if name != '_exists' and not super().__getattribute__('_exists'):
300 raise HandledException('Object does not exist.')
301 return super().__getattribute__(name)
303 def _disappear(self) -> None:
304 """Invalidate object, make future use raise exceptions."""
305 assert self.id_ is not None
306 if self._get_cached(self.id_):
308 to_kill = list(self.__dict__.keys())
314 def empty_cache(cls) -> None:
315 """Empty class's cache."""
319 def get_cache(cls: type[BaseModelInstance]) -> dict[Any, BaseModel[Any]]:
320 """Get cache dictionary, create it if not yet existing."""
321 if not hasattr(cls, 'cache_'):
322 d: dict[Any, BaseModel[Any]] = {}
327 def _get_cached(cls: type[BaseModelInstance],
328 id_: BaseModelId) -> BaseModelInstance | None:
329 """Get object of id_ from class's cache, or None if not found."""
330 # pylint: disable=consider-iterating-dictionary
331 cache = cls.get_cache()
332 if id_ in cache.keys():
334 assert isinstance(obj, cls)
338 def _cache(self) -> None:
339 """Update object in class's cache.
341 Also calls ._disappear if cache holds older reference to object of same
342 ID, but different memory address, to avoid doing anything with
346 raise HandledException('Cannot cache object without ID.')
347 cache = self.get_cache()
348 old_cached = self._get_cached(self.id_)
349 if old_cached and id(old_cached) != id(self):
350 # pylint: disable=protected-access
351 # (cause we remain within the class)
352 old_cached._disappear()
353 cache[self.id_] = self
355 def _uncache(self) -> None:
356 """Remove self from cache."""
358 raise HandledException('Cannot un-cache object without ID.')
359 cache = self.get_cache()
362 # object retrieval and generation
365 def from_table_row(cls: type[BaseModelInstance],
366 # pylint: disable=unused-argument
367 db_conn: DatabaseConnection,
368 row: Row | list[Any]) -> BaseModelInstance:
369 """Make from DB row, update DB cache with it."""
375 def by_id(cls, db_conn: DatabaseConnection,
376 id_: BaseModelId | None,
377 # pylint: disable=unused-argument
378 create: bool = False) -> Self:
379 """Retrieve by id_, on failure throw NotFoundException.
381 First try to get from cls.cache_, only then check DB; if found,
384 If create=True, make anew (but do not cache yet).
388 obj = cls._get_cached(id_)
390 for row in db_conn.row_where(cls.table_name, 'id', id_):
391 obj = cls.from_table_row(db_conn, row)
398 raise NotFoundException(f'found no object of ID {id_}')
401 def all(cls: type[BaseModelInstance],
402 db_conn: DatabaseConnection) -> list[BaseModelInstance]:
403 """Collect all objects of class into list.
405 Note that this primarily returns the contents of the cache, and only
406 _expands_ that by additional findings in the DB. This assumes the
407 cache is always instantly cleaned of any items that would be removed
410 items: dict[BaseModelId, BaseModelInstance] = {}
411 for k, v in cls.get_cache().items():
412 assert isinstance(v, cls)
414 already_recorded = items.keys()
415 for id_ in db_conn.column_all(cls.table_name, 'id'):
416 if id_ not in already_recorded:
417 item = cls.by_id(db_conn, id_)
418 assert item.id_ is not None
419 items[item.id_] = item
420 return list(items.values())
423 def by_date_range_with_limits(cls: type[BaseModelInstance],
424 db_conn: DatabaseConnection,
425 date_range: tuple[str, str],
426 date_col: str = 'day'
427 ) -> tuple[list[BaseModelInstance], str,
429 """Return list of items in database within (open) date_range interval.
431 If no range values provided, defaults them to 'yesterday' and
432 'tomorrow'. Knows to properly interpret these and 'today' as value.
434 start_str = date_range[0] if date_range[0] else 'yesterday'
435 end_str = date_range[1] if date_range[1] else 'tomorrow'
436 start_date = valid_date(start_str)
437 end_date = valid_date(end_str)
439 sql = f'SELECT id FROM {cls.table_name} '
440 sql += f'WHERE {date_col} >= ? AND {date_col} <= ?'
441 for row in db_conn.exec(sql, (start_date, end_date)):
442 items += [cls.by_id(db_conn, row[0])]
443 return items, start_date, end_date
446 def matching(cls: type[BaseModelInstance], db_conn: DatabaseConnection,
447 pattern: str) -> list[BaseModelInstance]:
448 """Return all objects whose .to_search match pattern."""
449 items = cls.all(db_conn)
453 for attr_name in cls.to_search:
454 toks = attr_name.split('.')
457 attr = getattr(parent, tok)
467 def save(self, db_conn: DatabaseConnection) -> None:
468 """Write self to DB and cache and ensure .id_.
470 Write both to DB, and to cache. To DB, write .id_ and attributes
471 listed in cls.to_save[_versioned|_relations].
473 Ensure self.id_ by setting it to what the DB command returns as the
474 last saved row's ID (cursor.lastrowid), EXCEPT if self.id_ already
475 exists as a 'str', which implies we do our own ID creation (so far
476 only the case with the Day class, where it's to be a date string.
478 values = tuple([self.id_] + [getattr(self, key)
479 for key in self.to_save])
480 table_name = self.table_name
481 cursor = db_conn.exec_on_vals(f'REPLACE INTO {table_name} VALUES',
483 if not isinstance(self.id_, str):
484 self.id_ = cursor.lastrowid # type: ignore[assignment]
486 for attr_name in self.to_save_versioned:
487 getattr(self, attr_name).save(db_conn)
488 for table, column, attr_name, key_index in self.to_save_relations:
489 assert isinstance(self.id_, (int, str))
490 db_conn.rewrite_relations(table, column, self.id_,
492 in getattr(self, attr_name)], key_index)
494 def remove(self, db_conn: DatabaseConnection) -> None:
495 """Remove from DB and cache, including dependencies."""
496 if self.id_ is None or self._get_cached(self.id_) is None:
497 raise HandledException('cannot remove unsaved item')
498 for attr_name in self.to_save_versioned:
499 getattr(self, attr_name).remove(db_conn)
500 for table, column, attr_name, _ in self.to_save_relations:
501 db_conn.delete_where(table, column, self.id_)
503 db_conn.delete_where(self.table_name, 'id', self.id_)