From 17f9fa203efb8d3bdad4ff3f1beeaed1381bd5fb Mon Sep 17 00:00:00 2001
From: John Lancaster <32917998+jsl12@users.noreply.github.com>
Date: Mon, 27 May 2024 00:08:38 -0500
Subject: [PATCH] more work

---
 README.md               | 11 +++++++++++
 requirements.txt        |  3 ++-
 src/restic/backup.py    | 11 +++++++----
 src/restic/loki.py      |  7 ++++++-
 src/restic/snapshots.py | 20 +++++++-------------
 5 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/README.md b/README.md
index 303d322..7116bdf 100644
--- a/README.md
+++ b/README.md
@@ -1,9 +1,20 @@
 # Restic Scripts
 
+## Environment Variables
+
+| Env Variable | Description                                                                                |
+|--------------|--------------------------------------------------------------------------------------------|
+| `BACKUP_DIR` | Directory to back up                                                                       |
+| `LOKI_URL`   | Push URL for Loki. Should include the port and end with something like `/loki/api/v1/push` |
+
 ## Usage
 
 ### Python
 
+```shell
+python -m restic.snapshots
+```
+
 ```shell
 python -m restic.backup
 ```
diff --git a/requirements.txt b/requirements.txt
index 749b7c7..f4c0c5b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 uv
 ruff
 rich
-requests
\ No newline at end of file
+requests
+click
\ No newline at end of file
diff --git a/src/restic/backup.py b/src/restic/backup.py
index 9eea237..47675f8 100755
--- a/src/restic/backup.py
+++ b/src/restic/backup.py
@@ -8,12 +8,13 @@ from rich.console import Console
 from rich.logging import RichHandler
 from rich.progress import Progress
 
+from restic import loki, snapshots
 from restic.console import console, logger
 from restic.loki import send_to_loki
 
 
-def run(backup_dir: Path, loki_url: str = None):
-    cmd = f'restic backup {backup_dir} --json --tag python-script'
+def run(backup_dir: Path, loki_url: str = None, tag: str = 'python-script'):
+    cmd = f'restic backup {backup_dir} --json --tag {tag}'
 
     logger.debug(f'Running cmd [bright_black]{cmd}[/]')
     process = subprocess.Popen(
@@ -54,14 +55,16 @@ def run(backup_dir: Path, loki_url: str = None):
     envvar='BACKUP_DIR',
 )
 @click.option('--loki-url', type=str, help='Loki URL for logging', envvar='LOKI_URL')
-def main(backup_dir: Path, loki_url: str = None):
+@click.option('--tag', type=str, help='Tag to use in restic')
+def main(backup_dir: Path, loki_url: str = None, tag: str = None):
     console = Console()
 
     logging.basicConfig(
         level='DEBUG', format='%(message)s', handlers=[RichHandler(markup=True, console=console)]
     )
 
-    run(backup_dir, loki_url)
+    run(backup_dir, loki_url, tag)
+    snapshots.run(loki_url)
 
 
 if __name__ == '__main__':
diff --git a/src/restic/loki.py b/src/restic/loki.py
index 9bcbb91..88c7a08 100644
--- a/src/restic/loki.py
+++ b/src/restic/loki.py
@@ -22,4 +22,9 @@ def send_to_loki(loki_url: str, line: str, backup: str):
     except Exception as e:
         logger.exception(e)
     else:
-        logger.info(f'Sent line to loki at {loki_url} {resp.text}')
+        if resp.status_code == 204:
+            logger.info(f'Sent line to loki at {loki_url}')
+        else:
+            logger.error(
+                f'Problem sending log line to Loki. Response status code {resp.status_code}: {resp.text}'
+            )
diff --git a/src/restic/snapshots.py b/src/restic/snapshots.py
index c51d90d..dc669b3 100644
--- a/src/restic/snapshots.py
+++ b/src/restic/snapshots.py
@@ -10,21 +10,15 @@ from restic.loki import send_to_loki
 
 
 def run(loki_url: str = None):
-    cmd = 'restic snapshots --json'
+    cmd = ['restic', 'snapshots', '--json']
 
-    logger.debug(f'Running cmd [bright_black]{cmd}[/]')
-    process = subprocess.Popen(
-        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True
-    )
+    with console.status('Getting snapshots...'):
+        logger.debug(f'Running cmd [bright_black]{" ".join(cmd)}[/]')
+        result = subprocess.run(cmd, capture_output=True, text=True)
 
-    with process.stdout:
-        for line in iter(process.stdout.readline, ''):
-            try:
-                data = json.loads(line)
-            except Exception as e:
-                logger.exception(e)
-            else:
-                console.print(data)
+    line = result.stdout
+    data = json.loads(line)
+    console.print(data)
 
     if loki_url is not None:
         send_to_loki(loki_url=loki_url, line=line, backup='snapshots')