Implement labels

- create label table
- create base database object class
- rename primary keys to 'id'
- log messages with label
- automatically reassign label to repo where possible

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

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

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):

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

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)

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
 

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):

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
 

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

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:

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

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

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

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,)

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:

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)