From 8792a34e8062f54ff8767d4901d187db000d4d2e Mon Sep 17 00:00:00 2001
From: John Lancaster <32917998+jsl12@users.noreply.github.com>
Date: Mon, 27 May 2024 12:49:20 -0500
Subject: [PATCH] added forget module

---
 src/restic/forget.py | 79 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)
 create mode 100644 src/restic/forget.py

diff --git a/src/restic/forget.py b/src/restic/forget.py
new file mode 100644
index 0000000..79d3b2d
--- /dev/null
+++ b/src/restic/forget.py
@@ -0,0 +1,79 @@
+import json
+import logging
+import subprocess
+from typing import Optional
+
+import click
+from pydantic import BaseModel
+from rich.logging import RichHandler
+
+from restic.console import console, logger
+from restic.loki import send_to_loki
+
+
+class KeepStrategy(BaseModel):
+    keep_last: int = 1
+    keep_hourly: Optional[int] = None
+    keep_daily: Optional[int] = None
+    keep_weekly: Optional[int] = None
+    keep_monthly: Optional[int] = None
+    keep_yearly: Optional[int] = None
+
+    def arguments(self):
+        for field in self.model_fields_set:
+            if val := getattr(self, field):
+                yield f'--{field.replace("_", "-")}', str(val)
+
+
+def forget(loki_url: str = None, dry_run: bool = False, **kwargs):
+    cmd = ['restic', 'forget', '--json', '--prune']
+
+    if dry_run:
+        cmd.append('--dry-run')
+
+    keep_strat = KeepStrategy.model_validate(kwargs)
+    for args in keep_strat.arguments():
+        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)
+        line = result.stdout
+        data = json.loads(line)
+    except json.JSONDecodeError:
+        logger.error('Error decoding JSON')
+        logger.error(line)
+    else:
+        if loki_url is not None and not dry_run:
+            send_to_loki(loki_url=loki_url, line=line, backup='forget')
+
+
+@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)
+@click.option('-l', '--keep-last', type=int, default=None)
+@click.option('-H', '--keep-hourly', type=int, default=None)
+@click.option('-d', '--keep-daily', type=int, default=None)
+@click.option('-w', '--keep-weekly', type=int, default=None)
+@click.option('-m', '--keep-monthly', type=int, default=None)
+@click.option('-y', '--keep-yearly', type=int, default=None)
+def main(loki_url: str, dry_run: bool, **kwargs):
+    logging.basicConfig(
+        level='DEBUG', format='%(message)s', handlers=[RichHandler(markup=True, console=console)]
+    )
+    logging.getLogger('docker.utils.config').setLevel('WARNING')
+    logging.getLogger('urllib3.connectionpool').setLevel('WARNING')
+
+    forget(loki_url, dry_run=dry_run, **kwargs)
+
+
+if __name__ == '__main__':
+    main()