From 57255d964cbaec1fb6cc38aed084e0c8d51bd4ac Mon Sep 17 00:00:00 2001
From: George Lacey <git@george.ooo>
Date: Tue, 4 May 2021 11:45:14 +0100
Subject: [PATCH] Handle borg output in separate class

---
 src/borgoutputhandler.py     | 24 ++++++++++++++++++++++++
 src/database/borgdatabase.py | 23 +++--------------------
 src/main.py                  | 28 ++++++++++++++++++++++------
 3 files changed, 49 insertions(+), 26 deletions(-)
 create mode 100644 src/borgoutputhandler.py

diff --git a/src/borgoutputhandler.py b/src/borgoutputhandler.py
new file mode 100644
index 0000000..6955cd4
--- /dev/null
+++ b/src/borgoutputhandler.py
@@ -0,0 +1,24 @@
+from datetime import datetime
+import borg
+import json
+
+
+class BorgOutputHandler(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)
+        except json.JSONDecodeError:
+            self.error = True
+
+    def get_borg_info(self):
+        repo = borg.Repo.from_json(self.borg_json['repository'])
+        archive = borg.Archive.from_json(self.borg_json['archive'])
+        stats = borg.Stats.from_json(self.borg_json['archive']['stats'])
+
+        return repo, archive, stats
+
+    def get_borg_error(self):
+        return borg.Error(self.borg_output, datetime.now())
diff --git a/src/database/borgdatabase.py b/src/database/borgdatabase.py
index 9b42653..7b004d6 100644
--- a/src/database/borgdatabase.py
+++ b/src/database/borgdatabase.py
@@ -1,8 +1,5 @@
 from .connection import RepoConn, ArchiveConn, StatsConn, ErrorConn
-from datetime import datetime
 from pathlib import Path
-from src import borg
-import json
 
 
 class BorgDatabase(object):
@@ -20,24 +17,10 @@ class BorgDatabase(object):
         self.error_conn = ErrorConn(db_path,
                                     table_name=self.error_name)
 
-    def process_borg_output(self, borg_output: str):
-        borg_json = None
-        try:
-            borg_json = json.loads(borg_output)
-        except json.JSONDecodeError:
-            self.handle_borg_error(borg_output)
-            return
-        self.process_borg_json(borg_json)
-
-    def process_borg_json(self, borg_json: dict):
-        repo = borg.Repo.from_json(borg_json['repository'])
-        archive = borg.Archive.from_json(borg_json['archive'])
-        stats = borg.Stats.from_json(borg_json['archive']['stats'])
-
+    def insert_record(self, repo, archive, stats):
         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)
 
-    def handle_borg_error(self, borg_error: str):
-        error = borg.Error(borg_error, datetime.now())
-        self.error_conn.insert(error)
+    def insert_error(self, borg_error):
+        self.error_conn.insert(borg_error)
diff --git a/src/main.py b/src/main.py
index c724d31..e9d7f61 100644
--- a/src/main.py
+++ b/src/main.py
@@ -2,16 +2,32 @@ from database import BorgDatabase
 from sys import stdin
 from os.path import realpath
 from pathlib import Path
+import argparse
+from borgoutputhandler import BorgOutputHandler
 
 
-def main(borg_output: str, path: Path):
-    db_path = path / 'borg.sqlite'
-    db = BorgDatabase(db_path)
+def main(args, path: Path):
+    if args.graph is not None:
+        pass
+    else:
+        borg_output = " ".join(stdin.readlines())
+        bo = BorgOutputHandler(borg_output)
+        db = BorgDatabase(path / 'borg.sqlite')
 
-    db.process_borg_output(borg_output)
+        if bo.error:
+            db.insert_error(bo.get_borg_error())
+        else:
+            db.insert_record(*bo.get_borg_info())
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-g", "--graph", help="Produce graphs at specified location",
+                        type=str)
+    return parser.parse_args()
 
 
 if __name__ == "__main__":
+    args = get_args()
     path = Path(realpath(__file__)).parent.parent
-    borg_output = "\n".join(stdin.readlines())
-    main(borg_output, path)
+    main(args, path)