fleshed out prune module

src/restic/forget.py

@@ -7,6 +7,7 @@ import click
 from pydantic import BaseModel
 from rich.logging import RichHandler
 
+from restic import size, snapshots
 from restic.console import console, logger
 from restic.loki import send_to_loki
 
@@ -26,7 +27,7 @@ class KeepStrategy(BaseModel):
 
 
 def forget(loki_url: str = None, dry_run: bool = False, **kwargs):
-    cmd = ['restic', 'forget', '--json', '--prune']
+    cmd = ['restic', 'forget', '--json']
 
     if dry_run:
         cmd.append('--dry-run')
@@ -36,10 +37,11 @@ def forget(loki_url: str = None, dry_run: bool = False, **kwargs):
         cmd.extend(args)
 
     logger.debug(f'Running cmd [bright_black]{" ".join(cmd)}[/]')
-
-    try:
     with console.status('Forgetting...'):
         result = subprocess.run(cmd, capture_output=True, text=True)
+
+    if result.returncode == 0:
+        try:
             line = result.stdout
             data = json.loads(line)
         except json.JSONDecodeError:
@@ -48,6 +50,11 @@ def forget(loki_url: str = None, dry_run: bool = False, **kwargs):
         else:
             if loki_url is not None and not dry_run:
                 send_to_loki(loki_url=loki_url, line=line, backup='forget')
+        finally:
+            logger.debug('Done')
+    else:
+        logger.error(f'Error running command')
+        logger.error(result.stderr)
 
 
 @click.command()
@@ -72,7 +79,13 @@ def main(loki_url: str, dry_run: bool, **kwargs):
     logging.getLogger('docker.utils.config').setLevel('WARNING')
     logging.getLogger('urllib3.connectionpool').setLevel('WARNING')
 
+    try:
         forget(loki_url, dry_run=dry_run, **kwargs)
+    except Exception:
+        console.print_exception()
+    else:
+        snapshots.snapshot(loki_url)
+        size.get_size(loki_url)
 
 
 if __name__ == '__main__':

src/restic/prune.py

@@ -4,8 +4,10 @@ import re
 import subprocess
 import sys
 
+import click
 from rich.logging import RichHandler
 
+from restic import size, snapshots
 from restic.console import console, logger
 from restic.loki import send_to_loki
 
@@ -52,12 +54,27 @@ def prune(loki_url: str = None, dry_run: bool = False):
         send_to_loki(loki_url, line=json.dumps(d), backup='prune')
 
 
-def main():
+@click.command()
+@click.option(
+    '--loki-url',
+    type=str,
+    help='Loki URL for logging. Defaults to the LOKI_URL env variable',
+    envvar='LOKI_URL',
+    default=None,
+)
+@click.option('-n', '--dry-run', type=bool, default=False, is_flag=True)
+def main(loki_url: str, dry_run: bool):
     logging.basicConfig(
         level='DEBUG', format='%(message)s', handlers=[RichHandler(markup=True, console=console)]
     )
 
-    prune(dry_run=True)
+    try:
+        prune(loki_url=loki_url, dry_run=dry_run)
+    except Exception as e:
+        raise e
+    else:
+        snapshots.snapshot(loki_url)
+        size.get_size(loki_url)
 
 
 if __name__ == '__main__':