diff --git a/src/borgmanager/borg/__init__.py b/src/borgmanager/borg/__init__.py index e4fc8cb..bc71aaa 100644 --- a/src/borgmanager/borg/__init__.py +++ b/src/borgmanager/borg/__init__.py @@ -1,5 +1,7 @@ +from .dbobject import DBObject from .repo import Repo from .archive import Archive from .stats import Stats from .error import Error +from .label import Label from .outputhandler import OutputHandler diff --git a/src/borgmanager/borg/dbobject.py b/src/borgmanager/borg/dbobject.py new file mode 100644 index 0000000..5987e2e --- /dev/null +++ b/src/borgmanager/borg/dbobject.py @@ -0,0 +1,13 @@ +from abc import ABC + + +class DBObject(ABC): + def __init__(self, primary_key=None): + self.__primary_key = primary_key + + @property + def primary_key(self): + if self.__primary_key is None: + raise ValueError("Primary key is None") + else: + return self.__primary_key diff --git a/src/borgmanager/borg/error.py b/src/borgmanager/borg/error.py index de1595a..0896482 100644 --- a/src/borgmanager/borg/error.py +++ b/src/borgmanager/borg/error.py @@ -1,11 +1,12 @@ +from . import DBObject from datetime import datetime -class Error(object): +class Error(DBObject): def __init__(self, error: str, time: datetime, primary_key=None): - self.error = error + super(Error, self).__init__(primary_key) + self.error = error.strip() self.time = time - self.primary_key = primary_key @classmethod def from_json(cls, json: dict): diff --git a/src/borgmanager/borg/label.py b/src/borgmanager/borg/label.py new file mode 100644 index 0000000..0a35dac --- /dev/null +++ b/src/borgmanager/borg/label.py @@ -0,0 +1,7 @@ +from . import DBObject + + +class Label(DBObject): + def __init__(self, label: str, primary_key=None): + super(Label, self).__init__(primary_key) + self.label = label diff --git a/src/borgmanager/borg/outputhandler.py b/src/borgmanager/borg/outputhandler.py index 98774f7..14dd799 100644 --- a/src/borgmanager/borg/outputhandler.py +++ b/src/borgmanager/borg/outputhandler.py @@ -7,6 +7,7 @@ class OutputHandler(object): def __init__(self, borg_output: str): self.borg_output = borg_output self.borg_json = None + self.error = False try: self.borg_json = json.loads(borg_output) diff --git a/src/borgmanager/borg/repo.py b/src/borgmanager/borg/repo.py index 33f91c0..c91e4df 100644 --- a/src/borgmanager/borg/repo.py +++ b/src/borgmanager/borg/repo.py @@ -1,20 +1,14 @@ +from . import DBObject from datetime import datetime from pathlib import Path -class Repo(object): +class Repo(DBObject): def __init__(self, fingerprint: str, location: Path, last_modified: datetime, primary_key=None): + super(Repo, self).__init__(primary_key) self.fingerprint = fingerprint self.location = location self.last_modified = last_modified - self.__primary_key = primary_key - - @property - def primary_key(self): - if self.__primary_key is None: - raise ValueError("Primary key is None") - else: - return self.__primary_key # region CLASS METHODS diff --git a/src/borgmanager/borg/stats.py b/src/borgmanager/borg/stats.py index 2c14284..94131d8 100644 --- a/src/borgmanager/borg/stats.py +++ b/src/borgmanager/borg/stats.py @@ -1,18 +1,14 @@ -class Stats(object): +from . import DBObject + + +class Stats(DBObject): def __init__(self, file_count: int, original_size: int, compressed_size: int, deduplicated_size: int, primary_key=None): + super(Stats, self).__init__(primary_key) self.file_count = file_count self.original_size = original_size self.compressed_size = compressed_size self.deduplicated_size = deduplicated_size - self.__primary_key = primary_key - - @property - def primary_key(self): - if self.__primary_key is None: - raise ValueError("Primary key is None") - else: - return self.__primary_key @classmethod def from_json(cls, json: dict): diff --git a/src/borgmanager/database/borgdatabase.py b/src/borgmanager/database/borgdatabase.py index 2d431fb..094acf6 100644 --- a/src/borgmanager/database/borgdatabase.py +++ b/src/borgmanager/database/borgdatabase.py @@ -1,4 +1,5 @@ -from .connection import RepoConn, ArchiveConn, StatsConn, ErrorConn +from .connection import RepoConn, ArchiveConn, StatsConn, ErrorConn, LabelConn +from borgmanager.borg.label import Label from pathlib import Path @@ -8,6 +9,7 @@ class BorgDatabase(object): self.archive_name = "archive" self.stats_name = "stats" self.error_name = "error" + self.label_name = "label" self.repo_conn = RepoConn(db_path, table_name=self.repo_name) self.archive_conn = ArchiveConn(db_path, self.repo_name, @@ -15,17 +17,26 @@ class BorgDatabase(object): self.stats_conn = StatsConn(db_path, self.repo_name, self.archive_name, table_name=self.stats_name) self.error_conn = ErrorConn(db_path, + label_table=self.label_name, table_name=self.error_name) + self.label_conn = LabelConn(db_path, + repo_table=self.repo_name, + table_name=self.label_name) # region INSERT - def insert_record(self, repo, archive, stats): + def insert_record(self, repo, archive, stats, label): repo_id = self.repo_conn.insert(repo) - archive_id = self.archive_conn.insert(archive, repo_id) - self.stats_conn.insert(stats, repo_id, archive_id) + label_id = self.insert_label(label, repo_id=repo_id) + archive_id = self.archive_conn.insert(archive, repo_id=repo_id) + self.stats_conn.insert(stats, repo_id=repo_id, archive_id=archive_id) - def insert_error(self, borg_error): - self.error_conn.insert(borg_error) + def insert_error(self, borg_error, label): + label_id = self.insert_label(label) + self.error_conn.insert(borg_error, label_id=label_id) + + def insert_label(self, label, repo_id=None): + return self.label_conn.insert(Label(label), repo_id=repo_id) # endregion diff --git a/src/borgmanager/database/connection/__init__.py b/src/borgmanager/database/connection/__init__.py index dc86103..9a32c7d 100644 --- a/src/borgmanager/database/connection/__init__.py +++ b/src/borgmanager/database/connection/__init__.py @@ -3,3 +3,4 @@ from .repoconn import RepoConn from .archiveconn import ArchiveConn from .statsconn import StatsConn from .errorconn import ErrorConn +from .labelconn import LabelConn diff --git a/src/borgmanager/database/connection/archiveconn.py b/src/borgmanager/database/connection/archiveconn.py index 6d14f41..e1e91ac 100644 --- a/src/borgmanager/database/connection/archiveconn.py +++ b/src/borgmanager/database/connection/archiveconn.py @@ -7,24 +7,23 @@ class ArchiveConn(DatabaseConnection): self.repo_table_name = repo_table_name super().__init__(db_path, table_name) - def _create_table(self): create_statement = f"create table if not exists {self._sql_table}(" \ - f"archive_id INTEGER PRIMARY KEY," \ + f"id INTEGER PRIMARY KEY," \ f"fingerprint TEXT NOT NULL UNIQUE," \ f"repo_id INTEGER NOT NULL," \ f"name TEXT NOT NULL," \ f"start TEXT TIMESTAMP NULL," \ f"end TEXT TIMESTAMP NULL," \ f"FOREIGN KEY (repo_id) REFERENCES" \ - f" {self.repo_table_name} (repo_id));" + f" {self.repo_table_name} (id));" self.sql_execute(create_statement) - def _exists(self, record): - return f"SELECT archive_id FROM {self._sql_table}" \ + def _exists(self, record, repo_id=None, archive_id=None, label_id=None): + return f"SELECT id FROM {self._sql_table}" \ f" WHERE fingerprint=?;", (record.fingerprint,) - def _insert(self, record, repo_id=None, archive_id=None) -> int: + def _insert(self, record, repo_id=None, archive_id=None, label_id=None) -> int: if repo_id is None: raise Exception("Repo id not supplied") with self.sql_lock: diff --git a/src/borgmanager/database/connection/databaseconnection.py b/src/borgmanager/database/connection/databaseconnection.py index acedc0b..82ec7c8 100644 --- a/src/borgmanager/database/connection/databaseconnection.py +++ b/src/borgmanager/database/connection/databaseconnection.py @@ -84,23 +84,23 @@ class DatabaseConnection(ABC): # region MODIFICATION - def insert(self, record, repo_id=None, archive_id=None): - exists, primary_key = self.exists(record) + def insert(self, record, repo_id=None, archive_id=None, label_id=None): + exists, primary_key = self.exists(record, repo_id, archive_id, label_id) if exists: - self._update(record, primary_key) + self._update(record, primary_key, repo_id, archive_id, label_id) return primary_key else: - return self._insert(record, repo_id, archive_id) + return self._insert(record, repo_id, archive_id, label_id) - def _update(self, record, primary_key): + def _update(self, record, primary_key, repo_id=None, archive_id=None, label_id=None): pass @abstractmethod - def _insert(self, record, repo_id=None, archive_id=None) -> int: + def _insert(self, record, repo_id=None, archive_id=None, label_id=None) -> int: raise NotImplementedError - def exists(self, record) -> (bool, int): - query, args = self._exists(record) + def exists(self, record, repo_id=None, archive_id=None, label_id=None) -> (bool, int): + query, args = self._exists(record, repo_id, archive_id, label_id) if query is None: return False, None @@ -112,7 +112,7 @@ class DatabaseConnection(ABC): return True, result[0] @abstractmethod - def _exists(self, record) -> (str, tuple): + def _exists(self, record, repo_id=None, archive_id=None, label_id=None) -> (str, tuple): raise NotImplementedError # endregion diff --git a/src/borgmanager/database/connection/errorconn.py b/src/borgmanager/database/connection/errorconn.py index f4c75ac..805cb4c 100644 --- a/src/borgmanager/database/connection/errorconn.py +++ b/src/borgmanager/database/connection/errorconn.py @@ -3,26 +3,32 @@ from datetime import datetime class ErrorConn(DatabaseConnection): - def __init__(self, db_path, table_name: str = "errors"): + def __init__(self, db_path, label_table: str, table_name: str = "errors"): + self.label_table = label_table super().__init__(db_path, table_name) def _create_table(self): create_statement = f"create table if not exists {self._sql_table}(" \ - f"error_id INTEGER PRIMARY KEY," \ + f"id INTEGER PRIMARY KEY," \ + f"label_id INT NOT NULL," \ f"error TEXT NOT NULL," \ - f"time TIMESTAMP NOT NULL);" + f"time TIMESTAMP NOT NULL," \ + f"FOREIGN KEY (label_id) REFERENCES" \ + f" {self.label_table} (id));" self.sql_execute(create_statement) - def _exists(self, record): + def _exists(self, record, repo_id=None, archive_id=None, label_id=None): return None, None - def _insert(self, record, repo_id=None, archive_id=None) -> int: + def _insert(self, record, repo_id=None, archive_id=None, label_id=None) -> int: + if label_id is None: + raise Exception("Label ID not supplied") with self.sql_lock: cursor = self.sql_cursor statement = f"INSERT INTO {self._sql_table}"\ - f" ('error', 'time')"\ - f" VALUES (?, ?);" - args = (record.error, datetime.now()) + f" ('label_id', 'error', 'time')"\ + f" VALUES (?, ?, ?);" + args = (label_id, record.error, datetime.now()) cursor.execute(statement, args) self.sql_commit() return cursor.lastrowid diff --git a/src/borgmanager/database/connection/labelconn.py b/src/borgmanager/database/connection/labelconn.py new file mode 100644 index 0000000..e11a370 --- /dev/null +++ b/src/borgmanager/database/connection/labelconn.py @@ -0,0 +1,59 @@ +from .databaseconnection import DatabaseConnection + + +class LabelConn(DatabaseConnection): + def __init__(self, db_path, repo_table: str, + table_name: str = "label"): + self.repo_table = repo_table + + super().__init__(db_path, table_name) + + # region INIT + + def _create_table(self): + create_statement = f"create table if not exists {self._sql_table}(" \ + f"id INTEGER PRIMARY KEY," \ + f"repo_id INT UNIQUE," \ + f"label TEXT NOT NULL," \ + f"FOREIGN KEY (repo_id) REFERENCES" \ + f" {self.repo_table} (id));" + self.sql_execute(create_statement) + + # endregion + + # region INSERT + + def _exists(self, record, repo_id=None, archive_id=None, label_id=None): + if repo_id is None: + return f"SELECT id FROM {self._sql_table} WHERE label=?;", (record.label,) + else: + return f"SELECT id FROM {self._sql_table} WHERE label=? OR repo_id=?;", (record.label, repo_id) + + def _insert(self, record, repo_id=None, archive_id=None, label_id=None) -> int: + with self.sql_lock: + cursor = self.sql_cursor + if repo_id is None: + statement = f"INSERT INTO {self._sql_table}"\ + f" ('label')"\ + f" VALUES (?);" + args = (record.label,) + else: + statement = f"INSERT INTO {self._sql_table}" \ + f" ('repo_id', 'label')" \ + f" VALUES (?, ?);" + args = (repo_id, record.label) + cursor.execute(statement, args) + self.sql_commit() + return cursor.lastrowid + + def _update(self, record, primary_key, repo_id=None, archive_id=None, label_id=None): + if repo_id is None: + self.sql_execute(f"UPDATE {self._sql_table} SET label = ? WHERE id = ?;", + (record.label, primary_key)) + else: + print("updating record") + self.sql_execute(f"UPDATE {self._sql_table} SET repo_id = ?, label = ? WHERE id = ?;", + (repo_id, record.label, primary_key)) + self.sql_commit() + + # endregion diff --git a/src/borgmanager/database/connection/repoconn.py b/src/borgmanager/database/connection/repoconn.py index ca83dc7..d078563 100644 --- a/src/borgmanager/database/connection/repoconn.py +++ b/src/borgmanager/database/connection/repoconn.py @@ -5,7 +5,15 @@ class RepoConn(DatabaseConnection): def __init__(self, db_path, table_name: str = 'repo'): super(RepoConn, self).__init__(db_path, table_name) - def _insert(self, record, repo_id=None, archive_id=None) -> int: + def _create_table(self): + create_statement = f"create table if not exists {self._sql_table}(" \ + f"id INTEGER PRIMARY KEY," \ + f"fingerprint TEXT NOT NULL UNIQUE," \ + f"location TEXT NOT NULL," \ + f"last_modified TIMESTAMP NOT NULL)" + self.sql_execute(create_statement) + + def _insert(self, record, repo_id=None, archive_id=None, label_id=None) -> int: with self.sql_lock: cursor = self.sql_cursor statement = f"INSERT INTO {self._sql_table}"\ @@ -16,18 +24,10 @@ class RepoConn(DatabaseConnection): self.sql_commit() return cursor.lastrowid - def _update(self, record, primary_key): - self.sql_execute(f"UPDATE {self._sql_table} SET location = ?, last_modified = ? WHERE repo_id = ?;", + def _update(self, record, primary_key, repo_id=None, archive_id=None, label_id=None): + self.sql_execute(f"UPDATE {self._sql_table} SET location = ?, last_modified = ? WHERE id = ?;", (str(record.location), record.last_modified, primary_key)) self.sql_commit() - def _exists(self, record): - return f"SELECT repo_id FROM {self._sql_table} WHERE fingerprint=?;", (record.fingerprint,) - - def _create_table(self): - create_statement = f"create table if not exists {self._sql_table}(" \ - f"repo_id INTEGER PRIMARY KEY," \ - f"fingerprint TEXT NOT NULL UNIQUE," \ - f"location TEXT NOT NULL," \ - f"last_modified TIMESTAMP NOT NULL)" - self.sql_execute(create_statement) + def _exists(self, record, repo_id=None, archive_id=None, label_id=None): + return f"SELECT id FROM {self._sql_table} WHERE fingerprint=?;", (record.fingerprint,) diff --git a/src/borgmanager/database/connection/statsconn.py b/src/borgmanager/database/connection/statsconn.py index b8fb93b..57f1239 100644 --- a/src/borgmanager/database/connection/statsconn.py +++ b/src/borgmanager/database/connection/statsconn.py @@ -13,7 +13,7 @@ class StatsConn(DatabaseConnection): def _create_table(self): create_statement = f"create table if not exists {self._sql_table}(" \ - f"stat_id INTEGER PRIMARY KEY," \ + f"id INTEGER PRIMARY KEY," \ f"repo_id INTEGER NOT NULL," \ f"archive_id INTEGER NOT NULL," \ f"file_count INTEGER NOT NULL," \ @@ -21,19 +21,19 @@ class StatsConn(DatabaseConnection): f"compressed_size INTEGER NOT NULL," \ f"deduplicated_size INTEGER NOT NULL," \ f"FOREIGN KEY (repo_id) REFERENCES" \ - f" {self.repo_table} (repo_id)," \ + f" {self.repo_table} (id)," \ f"FOREIGN KEY (archive_id) REFERENCES" \ - f" {self.archive_table} (archive_id));" + f" {self.archive_table} (id));" self.sql_execute(create_statement) # endregion # region INSERT - def _exists(self, record): + def _exists(self, record, repo_id=None, archive_id=None, label_id=None): return None, None - def _insert(self, record, repo_id=None, archive_id=None) -> int: + def _insert(self, record, repo_id=None, archive_id=None, label_id=None) -> int: if repo_id is None or archive_id is None: raise Exception("Repo and archive ids not supplied") with self.sql_lock: diff --git a/src/main.py b/src/main.py index 2d47dec..3554ff9 100644 --- a/src/main.py +++ b/src/main.py @@ -21,12 +21,15 @@ def main(args, path: Path): summary = Summary(db, args.summary) else: borg_output = " ".join(stdin.readlines()) - bo = OutputHandler(borg_output) - - if bo.error: - db.insert_error(bo.get_borg_error()) + if args.label is None: + raise Exception("No label supplied") else: - db.insert_record(*bo.get_borg_info()) + bo = OutputHandler(borg_output) + + if bo.error: + db.insert_error(bo.get_borg_error(), args.label) + else: + db.insert_record(*bo.get_borg_info(), args.label) def get_args(): @@ -34,10 +37,11 @@ def get_args(): parser.add_argument("-g", "--graph", help="Produce graphs at specified location", type=str) parser.add_argument("-s", "--summary", help="Print summary", type=str) parser.add_argument("-o", "--output", help="Output Directory", type=str) + parser.add_argument("-l", "--label", help="Repo Label", type=str) return parser.parse_args() if __name__ == "__main__": - args = get_args() - path = Path(realpath(__file__)).parent.parent - main(args, path) + m_args = get_args() + m_path = Path(realpath(__file__)).parent.parent + main(m_args, m_path)