plomlompom.com Git - plomtask/blob - plomtask/db.py

   1 """Database management."""
   2 from __future__ import annotations
   3 from os import listdir
   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
  10 EXPECTED_DB_VERSION = 1
  11 MIGRATIONS_DIR = 'migrations'
  12 FILENAME_DB_SCHEMA = f'init_{EXPECTED_DB_VERSION}.sql'
  13 PATH_DB_SCHEMA = f'{MIGRATIONS_DIR}/{FILENAME_DB_SCHEMA}'
  14
  15
  16 class UnmigratedDbException(HandledException):
  17     """To identify case of unmigrated DB file."""
  18
  19
  20 class DatabaseFile:  # pylint: disable=too-few-public-methods
  21     """Represents the sqlite3 database's file."""
  22
  23     def __init__(self, path: str) -> None:
  24         self.path = path
  25         self._check()
  26
  27     @classmethod
  28     def create_at(cls, path: str) -> DatabaseFile:
  29         """Make new DB file at path."""
  30         with sql_connect(path) as conn:
  31             with open(PATH_DB_SCHEMA, 'r', encoding='utf-8') as f:
  32                 conn.executescript(f.read())
  33             conn.execute(f'PRAGMA user_version = {EXPECTED_DB_VERSION}')
  34         return cls(path)
  35
  36     @classmethod
  37     def migrate(cls, path: str) -> DatabaseFile:
  38         """Apply migrations from_version to EXPECTED_DB_VERSION."""
  39         migrations = cls._available_migrations()
  40         from_version = cls.get_version_of_db(path)
  41         migrations_todo = migrations[from_version+1:]
  42         for j, filename in enumerate(migrations_todo):
  43             with sql_connect(path) as conn:
  44                 with open(f'{MIGRATIONS_DIR}/{filename}', 'r',
  45                           encoding='utf-8') as f:
  46                     conn.executescript(f.read())
  47             user_version = from_version + j + 1
  48             with sql_connect(path) as conn:
  49                 conn.execute(f'PRAGMA user_version = {user_version}')
  50         return cls(path)
  51
  52     def _check(self) -> None:
  53         """Check file exists, and is of proper DB version and schema."""
  54         if not isfile(self.path):
  55             raise NotFoundException
  56         if self.user_version != EXPECTED_DB_VERSION:
  57             raise UnmigratedDbException()
  58         self._validate_schema()
  59
  60     @staticmethod
  61     def _available_migrations() -> list[str]:
  62         """Validate migrations directory and return sorted entries."""
  63         msg_too_big = 'Migration directory points beyond expected DB version.'
  64         msg_bad_entry = 'Migration directory contains unexpected entry: '
  65         msg_missing = 'Migration directory misses migration of number: '
  66         migrations = {}
  67         for entry in listdir(MIGRATIONS_DIR):
  68             if entry == FILENAME_DB_SCHEMA:
  69                 continue
  70             toks = entry.split('_', 1)
  71             if len(toks) < 2:
  72                 raise HandledException(msg_bad_entry + entry)
  73             try:
  74                 i = int(toks[0])
  75             except ValueError as e:
  76                 raise HandledException(msg_bad_entry + entry) from e
  77             if i > EXPECTED_DB_VERSION:
  78                 raise HandledException(msg_too_big)
  79             migrations[i] = toks[1]
  80         migrations_list = []
  81         for i in range(EXPECTED_DB_VERSION + 1):
  82             if i not in migrations:
  83                 raise HandledException(msg_missing + str(i))
  84             migrations_list += [f'{i}_{migrations[i]}']
  85         return migrations_list
  86
  87     @staticmethod
  88     def get_version_of_db(path: str) -> int:
  89         """Get DB user_version, fail if outside expected range."""
  90         sql_for_db_version = 'PRAGMA user_version'
  91         with sql_connect(path) as conn:
  92             db_version = list(conn.execute(sql_for_db_version))[0][0]
  93         if db_version > EXPECTED_DB_VERSION:
  94             msg = f'Wrong DB version, expected '\
  95                     f'{EXPECTED_DB_VERSION}, got unknown {db_version}.'
  96             raise HandledException(msg)
  97         assert isinstance(db_version, int)
  98         return db_version
  99
 100     @property
 101     def user_version(self) -> int:
 102         """Get DB user_version."""
 103         return self.__class__.get_version_of_db(self.path)
 104
 105     def _validate_schema(self) -> None:
 106         """Compare found schema with what's stored at PATH_DB_SCHEMA."""
 107
 108         def reformat_rows(rows: list[str]) -> list[str]:
 109             new_rows = []
 110             for row in rows:
 111                 new_row = []
 112                 for subrow in row.split('\n'):
 113                     subrow = subrow.rstrip()
 114                     in_parentheses = 0
 115                     split_at = []
 116                     for i, c in enumerate(subrow):
 117                         if '(' == c:
 118                             in_parentheses += 1
 119                         elif ')' == c:
 120                             in_parentheses -= 1
 121                         elif ',' == c and 0 == in_parentheses:
 122                             split_at += [i + 1]
 123                     prev_split = 0
 124                     for i in split_at:
 125                         segment = subrow[prev_split:i].strip()
 126                         if len(segment) > 0:
 127                             new_row += [f'    {segment}']
 128                         prev_split = i
 129                     segment = subrow[prev_split:].strip()
 130                     if len(segment) > 0:
 131                         new_row += [f'    {segment}']
 132                 new_row[0] = new_row[0].lstrip()
 133                 new_row[-1] = new_row[-1].lstrip()
 134                 new_rows += ['\n'.join(new_row)]
 135             return new_rows
 136
 137         sql_for_schema = 'SELECT sql FROM sqlite_master ORDER BY sql'
 138         msg_err = 'Database has wrong tables schema. Diff:\n'
 139         with sql_connect(self.path) as conn:
 140             schema_rows = [r[0] for r in conn.execute(sql_for_schema) if r[0]]
 141         schema_rows = reformat_rows(schema_rows)
 142         retrieved_schema = ';\n'.join(schema_rows) + ';'
 143         with open(PATH_DB_SCHEMA, 'r', encoding='utf-8') as f:
 144             stored_schema = f.read().rstrip()
 145         if stored_schema != retrieved_schema:
 146             diff_msg = Differ().compare(retrieved_schema.splitlines(),
 147                                         stored_schema.splitlines())
 148             raise HandledException(msg_err + '\n'.join(diff_msg))
 149
 150
 151 class DatabaseConnection:
 152     """A single connection to the database."""
 153
 154     def __init__(self, db_file: DatabaseFile) -> None:
 155         self.file = db_file
 156         self.conn = sql_connect(self.file.path)
 157
 158     def commit(self) -> None:
 159         """Commit SQL transaction."""
 160         self.conn.commit()
 161
 162     def exec(self, code: str, inputs: tuple[Any, ...] = tuple()) -> Cursor:
 163         """Add commands to SQL transaction."""
 164         return self.conn.execute(code, inputs)
 165
 166     def close(self) -> None:
 167         """Close DB connection."""
 168         self.conn.close()
 169
 170     def rewrite_relations(self, table_name: str, key: str, target: int | str,
 171                           rows: list[list[Any]]) -> None:
 172         """Rewrite relations in table_name to target, with rows values."""
 173         self.delete_where(table_name, key, target)
 174         for row in rows:
 175             values = tuple([target] + row)
 176             q_marks = self.__class__.q_marks_from_values(values)
 177             self.exec(f'INSERT INTO {table_name} VALUES {q_marks}', values)
 178
 179     def row_where(self, table_name: str, key: str,
 180                   target: int | str) -> list[Row]:
 181         """Return list of Rows at table where key == target."""
 182         return list(self.exec(f'SELECT * FROM {table_name} WHERE {key} = ?',
 183                               (target,)))
 184
 185     def column_where(self, table_name: str, column: str, key: str,
 186                      target: int | str) -> list[Any]:
 187         """Return column of table where key == target."""
 188         return [row[0] for row in
 189                 self.exec(f'SELECT {column} FROM {table_name} '
 190                           f'WHERE {key} = ?', (target,))]
 191
 192     def column_all(self, table_name: str, column: str) -> list[Any]:
 193         """Return complete column of table."""
 194         return [row[0] for row in
 195                 self.exec(f'SELECT {column} FROM {table_name}')]
 196
 197     def delete_where(self, table_name: str, key: str,
 198                      target: int | str) -> None:
 199         """Delete from table where key == target."""
 200         self.exec(f'DELETE FROM {table_name} WHERE {key} = ?', (target,))
 201
 202     @staticmethod
 203     def q_marks_from_values(values: tuple[Any]) -> str:
 204         """Return placeholder to insert values into SQL code."""
 205         return '(' + ','.join(['?'] * len(values)) + ')'
 206
 207
 208 BaseModelId = TypeVar('BaseModelId', int, str)
 209 BaseModelInstance = TypeVar('BaseModelInstance', bound='BaseModel[Any]')
 210
 211
 212 class BaseModel(Generic[BaseModelId]):
 213     """Template for most of the models we use/derive from the DB."""
 214     table_name = ''
 215     to_save: list[str] = []
 216     to_save_versioned: list[str] = []
 217     to_save_relations: list[tuple[str, str, str]] = []
 218     id_: None | BaseModelId
 219     cache_: dict[BaseModelId, Self]
 220
 221     def __init__(self, id_: BaseModelId | None) -> None:
 222         if isinstance(id_, int) and id_ < 1:
 223             msg = f'illegal {self.__class__.__name__} ID, must be >=1: {id_}'
 224             raise HandledException(msg)
 225         self.id_ = id_
 226
 227     def __eq__(self, other: object) -> bool:
 228         if not isinstance(other, self.__class__):
 229             return False
 230         to_hash_me = tuple([self.id_] +
 231                            [getattr(self, name) for name in self.to_save])
 232         to_hash_other = tuple([other.id_] +
 233                               [getattr(other, name) for name in other.to_save])
 234         return hash(to_hash_me) == hash(to_hash_other)
 235
 236     def __lt__(self, other: Any) -> bool:
 237         if not isinstance(other, self.__class__):
 238             msg = 'cannot compare to object of different class'
 239             raise HandledException(msg)
 240         assert isinstance(self.id_, int)
 241         assert isinstance(other.id_, int)
 242         return self.id_ < other.id_
 243
 244     @classmethod
 245     def get_cached(cls: type[BaseModelInstance],
 246                    id_: BaseModelId) -> BaseModelInstance | None:
 247         """Get object of id_ from class's cache, or None if not found."""
 248         # pylint: disable=consider-iterating-dictionary
 249         cache = cls.get_cache()
 250         if id_ in cache.keys():
 251             obj = cache[id_]
 252             assert isinstance(obj, cls)
 253             return obj
 254         return None
 255
 256     @classmethod
 257     def empty_cache(cls) -> None:
 258         """Empty class's cache."""
 259         cls.cache_ = {}
 260
 261     @classmethod
 262     def get_cache(cls: type[BaseModelInstance]) -> dict[Any, BaseModel[Any]]:
 263         """Get cache dictionary, create it if not yet existing."""
 264         if not hasattr(cls, 'cache_'):
 265             d: dict[Any, BaseModel[Any]] = {}
 266             cls.cache_ = d
 267         return cls.cache_
 268
 269     def cache(self) -> None:
 270         """Update object in class's cache."""
 271         if self.id_ is None:
 272             raise HandledException('Cannot cache object without ID.')
 273         cache = self.__class__.get_cache()
 274         cache[self.id_] = self
 275
 276     def uncache(self) -> None:
 277         """Remove self from cache."""
 278         if self.id_ is None:
 279             raise HandledException('Cannot un-cache object without ID.')
 280         cache = self.__class__.get_cache()
 281         del cache[self.id_]
 282
 283     @classmethod
 284     def from_table_row(cls: type[BaseModelInstance],
 285                        # pylint: disable=unused-argument
 286                        db_conn: DatabaseConnection,
 287                        row: Row | list[Any]) -> BaseModelInstance:
 288         """Make from DB row, write to DB cache."""
 289         obj = cls(*row)
 290         obj.cache()
 291         return obj
 292
 293     @classmethod
 294     def by_id(cls, db_conn: DatabaseConnection,
 295               id_: BaseModelId | None,
 296               # pylint: disable=unused-argument
 297               create: bool = False) -> Self:
 298         """Retrieve by id_, on failure throw NotFoundException.
 299
 300         First try to get from cls.cache_, only then check DB; if found,
 301         put into cache.
 302
 303         If create=True, make anew (but do not cache yet).
 304         """
 305         obj = None
 306         if id_ is not None:
 307             obj = cls.get_cached(id_)
 308             if not obj:
 309                 for row in db_conn.row_where(cls.table_name, 'id', id_):
 310                     obj = cls.from_table_row(db_conn, row)
 311                     obj.cache()
 312                     break
 313         if obj:
 314             return obj
 315         if create:
 316             obj = cls(id_)
 317             return obj
 318         raise NotFoundException(f'found no object of ID {id_}')
 319
 320     @classmethod
 321     def all(cls: type[BaseModelInstance],
 322             db_conn: DatabaseConnection) -> list[BaseModelInstance]:
 323         """Collect all objects of class into list.
 324
 325         Note that this primarily returns the contents of the cache, and only
 326         _expands_ that by additional findings in the DB. This assumes the
 327         cache is always instantly cleaned of any items that would be removed
 328         from the DB.
 329         """
 330         items: dict[BaseModelId, BaseModelInstance] = {}
 331         for k, v in cls.get_cache().items():
 332             assert isinstance(v, cls)
 333             items[k] = v
 334         already_recorded = items.keys()
 335         for id_ in db_conn.column_all(cls.table_name, 'id'):
 336             if id_ not in already_recorded:
 337                 item = cls.by_id(db_conn, id_)
 338                 assert item.id_ is not None
 339                 items[item.id_] = item
 340         return list(items.values())
 341
 342     def save(self, db_conn: DatabaseConnection) -> None:
 343         """Write self to DB and cache and ensure .id_.
 344
 345         Write both to DB, and to cache. To DB, write .id_ and attributes
 346         listed in cls.to_save[_versioned|_relations].
 347
 348         Ensure self.id_ by setting it to what the DB command returns as the
 349         last saved row's ID (cursor.lastrowid), EXCEPT if self.id_ already
 350         exists as a 'str', which implies we do our own ID creation (so far
 351         only the case with the Day class, where it's to be a date string.
 352         """
 353         values = tuple([self.id_] + [getattr(self, key)
 354                                      for key in self.to_save])
 355         q_marks = DatabaseConnection.q_marks_from_values(values)
 356         table_name = self.table_name
 357         cursor = db_conn.exec(f'REPLACE INTO {table_name} VALUES {q_marks}',
 358                               values)
 359         if not isinstance(self.id_, str):
 360             self.id_ = cursor.lastrowid  # type: ignore[assignment]
 361         self.cache()
 362         for attr_name in self.to_save_versioned:
 363             getattr(self, attr_name).save(db_conn)
 364         for table, column, attr_name in self.to_save_relations:
 365             assert isinstance(self.id_, (int, str))
 366             db_conn.rewrite_relations(table, column, self.id_,
 367                                       [[i.id_] for i
 368                                        in getattr(self, attr_name)])
 369
 370     def remove(self, db_conn: DatabaseConnection) -> None:
 371         """Remove from DB and cache, including dependencies."""
 372         if self.id_ is None or self.__class__.get_cached(self.id_) is None:
 373             raise HandledException('cannot remove unsaved item')
 374         for attr_name in self.to_save_versioned:
 375             getattr(self, attr_name).remove(db_conn)
 376         for table, column, attr_name in self.to_save_relations:
 377             db_conn.delete_where(table, column, self.id_)
 378         self.uncache()
 379         db_conn.delete_where(self.table_name, 'id', self.id_)