home · contact · privacy
dfb388b7ec6ed9317f19a788336025a96e4c745e
[plomtask] / plomtask / db.py
1 """Database management."""
2 from os.path import isfile
3 from difflib import Differ
4 from sqlite3 import connect as sql_connect, Cursor
5 from typing import Any, Dict
6 from plomtask.exceptions import HandledException
7
8 PATH_DB_SCHEMA = 'scripts/init.sql'
9
10
11 class DatabaseFile:  # pylint: disable=too-few-public-methods
12     """Represents the sqlite3 database's file."""
13
14     def __init__(self, path: str) -> None:
15         self.path = path
16         self._check()
17
18     def remake(self) -> None:
19         """Create tables in self.path file as per PATH_DB_SCHEMA sql file."""
20         with sql_connect(self.path) as conn:
21             with open(PATH_DB_SCHEMA, 'r', encoding='utf-8') as f:
22                 conn.executescript(f.read())
23         self._check()
24
25     def _check(self) -> None:
26         """Check file exists, and is of proper schema."""
27         self.exists = isfile(self.path)
28         if self.exists:
29             self._validate_schema()
30
31     def _validate_schema(self) -> None:
32         """Compare found schema with what's stored at PATH_DB_SCHEMA."""
33         sql_for_schema = 'SELECT sql FROM sqlite_master ORDER BY sql'
34         msg_err = 'Database has wrong tables schema. Diff:\n'
35         with sql_connect(self.path) as conn:
36             schema_rows = [r[0] for r in conn.execute(sql_for_schema) if r[0]]
37             retrieved_schema = ';\n'.join(schema_rows) + ';'
38             with open(PATH_DB_SCHEMA, 'r', encoding='utf-8') as f:
39                 stored_schema = f.read().rstrip()
40                 if stored_schema != retrieved_schema:
41                     diff_msg = Differ().compare(retrieved_schema.splitlines(),
42                                                 stored_schema.splitlines())
43                     raise HandledException(msg_err + '\n'.join(diff_msg))
44
45
46 class DatabaseConnection:
47     """A single connection to the database."""
48
49     def __init__(self, db_file: DatabaseFile) -> None:
50         self.file = db_file
51         self.conn = sql_connect(self.file.path)
52         self.cached_todos: Dict[int, Any] = {}
53         self.cached_days: Dict[str, Any] = {}
54         self.cached_process_steps: Dict[int, Any] = {}
55         self.cached_processes: Dict[int, Any] = {}
56         self.cached_conditions: Dict[int, Any] = {}
57
58     def commit(self) -> None:
59         """Commit SQL transaction."""
60         self.conn.commit()
61
62     def exec(self, code: str, inputs: tuple[Any, ...] = tuple()) -> Cursor:
63         """Add commands to SQL transaction."""
64         return self.conn.execute(code, inputs)
65
66     def close(self) -> None:
67         """Close DB connection."""
68         self.conn.close()
69
70
71 class BaseModel:
72     """Template for most of the models we use/derive from the DB."""
73     table_name = ''
74     to_save: list[str] = []
75     id_: None | int | str
76
77     def set_int_id(self, id_: int | None) -> None:
78         """Set id_ if >= 1 or None, else fail."""
79         if (id_ is not None) and id_ < 1:
80             msg = f'illegal {self.__class__.__name__} ID, must be >=1: {id_}'
81             raise HandledException(msg)
82         self.id_ = id_
83
84     def save_core(self, db_conn: DatabaseConnection,
85                   update_with_lastrowid: bool = True) -> None:
86         """Write bare-bones self (sans connected items), ensuring self.id_."""
87         q_marks = ','.join(['?'] * (len(self.to_save) + 1))
88         values = tuple([self.id_] + [getattr(self, key)
89                                      for key in self.to_save])
90         table_name = self.table_name
91         cursor = db_conn.exec(f'REPLACE INTO {table_name} VALUES ({q_marks})',
92                               values)
93         if update_with_lastrowid:
94             self.id_ = cursor.lastrowid
95         cache = getattr(db_conn, f'cached_{table_name}')
96         cache[self.id_] = self