X-Git-Url: https://plomlompom.com/repos/berlin_corona.txt?a=blobdiff_plain;f=plomtask%2Fdb.py;h=bf2b79f758b313525816581620a792330fdbd9a2;hb=23c7ff7f9833ff5b0e547c19a4ad85325a5d3400;hp=50eb737cbfa3924eb3f393dd5440ed08fce0b8d9;hpb=9c3ba57169fdc71d3ffefd4c7c20cae4e89f9c9c;p=plomtask diff --git a/plomtask/db.py b/plomtask/db.py index 50eb737..bf2b79f 100644 --- a/plomtask/db.py +++ b/plomtask/db.py @@ -1,33 +1,47 @@ """Database management.""" +from __future__ import annotations from os.path import isfile from difflib import Differ -from sqlite3 import connect as sql_connect -from plomtask.misc import HandledException +from sqlite3 import connect as sql_connect, Cursor, Row +from typing import Any, Self, TypeVar, Generic +from plomtask.exceptions import HandledException PATH_DB_SCHEMA = 'scripts/init.sql' +EXPECTED_DB_VERSION = 0 -class DatabaseFile: +class DatabaseFile: # pylint: disable=too-few-public-methods """Represents the sqlite3 database's file.""" - def __init__(self, path): + def __init__(self, path: str) -> None: self.path = path - self.check() + self._check() - def check(self): - """Check file exists and is of proper schema.""" - self.exists = isfile(self.path) - if self.exists: - self.validate_schema() - - def remake(self): + def remake(self) -> None: """Create tables in self.path file as per PATH_DB_SCHEMA sql file.""" with sql_connect(self.path) as conn: with open(PATH_DB_SCHEMA, 'r', encoding='utf-8') as f: conn.executescript(f.read()) - self.check() + self._check() - def validate_schema(self): + def _check(self) -> None: + """Check file exists, and is of proper DB version and schema.""" + self.exists = isfile(self.path) + if self.exists: + self._validate_user_version() + self._validate_schema() + + def _validate_user_version(self) -> None: + """Compare DB user_version with EXPECTED_DB_VERSION.""" + sql_for_db_version = 'PRAGMA user_version' + with sql_connect(self.path) as conn: + db_version = list(conn.execute(sql_for_db_version))[0][0] + if db_version != EXPECTED_DB_VERSION: + msg = f'Wrong DB version, expected '\ + f'{EXPECTED_DB_VERSION}, got {db_version}.' + raise HandledException(msg) + + def _validate_schema(self) -> None: """Compare found schema with what's stored at PATH_DB_SCHEMA.""" sql_for_schema = 'SELECT sql FROM sqlite_master ORDER BY sql' msg_err = 'Database has wrong tables schema. Diff:\n' @@ -45,18 +59,139 @@ class DatabaseFile: class DatabaseConnection: """A single connection to the database.""" - def __init__(self, db_file: DatabaseFile): + def __init__(self, db_file: DatabaseFile) -> None: self.file = db_file self.conn = sql_connect(self.file.path) - def commit(self): + def commit(self) -> None: """Commit SQL transaction.""" self.conn.commit() - def exec(self, code: str, inputs: tuple = ()): + def exec(self, code: str, inputs: tuple[Any, ...] = tuple()) -> Cursor: """Add commands to SQL transaction.""" return self.conn.execute(code, inputs) - def close(self): + def close(self) -> None: """Close DB connection.""" self.conn.close() + + def rewrite_relations(self, table_name: str, key: str, target: int, + rows: list[list[Any]]) -> None: + """Rewrite relations in table_name to target, with rows values.""" + self.delete_where(table_name, key, target) + for row in rows: + values = tuple([target] + row) + q_marks = self.__class__.q_marks_from_values(values) + self.exec(f'INSERT INTO {table_name} VALUES {q_marks}', values) + + def row_where(self, table_name: str, key: str, + target: int | str) -> list[Row]: + """Return list of Rows at table where key == target.""" + return list(self.exec(f'SELECT * FROM {table_name} WHERE {key} = ?', + (target,))) + + def column_where(self, table_name: str, column: str, key: str, + target: int | str) -> list[Any]: + """Return column of table where key == target.""" + return [row[0] for row in + self.exec(f'SELECT {column} FROM {table_name} ' + f'WHERE {key} = ?', (target,))] + + def column_all(self, table_name: str, column: str) -> list[Any]: + """Return complete column of table.""" + return [row[0] for row in + self.exec(f'SELECT {column} FROM {table_name}')] + + def delete_where(self, table_name: str, key: str, target: int) -> None: + """Delete from table where key == target.""" + self.exec(f'DELETE FROM {table_name} WHERE {key} = ?', (target,)) + + @staticmethod + def q_marks_from_values(values: tuple[Any]) -> str: + """Return placeholder to insert values into SQL code.""" + return '(' + ','.join(['?'] * len(values)) + ')' + + +X = TypeVar('X', int, str) +T = TypeVar('T', bound='BaseModel[Any]') + + +class BaseModel(Generic[X]): + """Template for most of the models we use/derive from the DB.""" + table_name = '' + to_save: list[str] = [] + id_: None | X + cache_: dict[X, Self] = {} + + @classmethod + def from_table_row(cls: type[T], + # pylint: disable=unused-argument + db_conn: DatabaseConnection, + row: Row | list[Any]) -> T: + """Make from DB row, write to DB cache.""" + obj = cls(*row) + obj.cache() + return obj + + @classmethod + def _by_id(cls, + db_conn: DatabaseConnection, + id_: X) -> tuple[Self | None, bool]: + """Return instance found by ID, or None, and if from cache or not.""" + from_cache = False + obj = cls.get_cached(id_) + if obj: + from_cache = True + else: + for row in db_conn.row_where(cls.table_name, 'id', id_): + obj = cls.from_table_row(db_conn, row) + obj.cache() + break + return obj, from_cache + + def set_int_id(self, id_: int | None) -> None: + """Set id_ if >= 1 or None, else fail.""" + if (id_ is not None) and id_ < 1: + msg = f'illegal {self.__class__.__name__} ID, must be >=1: {id_}' + raise HandledException(msg) + self.id_ = id_ # type: ignore[assignment] + + def save_core(self, db_conn: DatabaseConnection, + update_with_lastrowid: bool = True) -> None: + """Write bare-bones self (sans connected items), ensuring self.id_.""" + values = tuple([self.id_] + [getattr(self, key) + for key in self.to_save]) + q_marks = DatabaseConnection.q_marks_from_values(values) + table_name = self.table_name + cursor = db_conn.exec(f'REPLACE INTO {table_name} VALUES {q_marks}', + values) + if update_with_lastrowid: + self.id_ = cursor.lastrowid # type: ignore[assignment] + self.cache() + + @classmethod + def get_cached(cls: type[T], id_: X) -> T | None: + """Get object of id_ from class's cache, or None if not found.""" + # pylint: disable=consider-iterating-dictionary + if id_ in cls.cache_.keys(): + obj = cls.cache_[id_] + assert isinstance(obj, cls) + return obj + return None + + def cache(self) -> None: + """Update object in class's cache.""" + if self.id_ is None: + raise HandledException('Cannot cache object without ID.') + self.__class__.cache_[self.id_] = self + + def uncache(self) -> None: + """Remove self from cache.""" + if self.id_ is None: + raise HandledException('Cannot un-cache object without ID.') + del self.__class__.cache_[self.id_] + + @classmethod + def empty_cache(cls) -> None: + """Empty class's cache.""" + cls.cache_ = {}