84 lines
2.1 KiB
Python
84 lines
2.1 KiB
Python
import json
|
|
import logging
|
|
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
|
|
|
|
field_regex = re.compile(
|
|
r'^(?P<name>[\w ]+):\s+(?P<blobs>\d+) blobs \/ (?P<size>\d+(\.\d+)? (?:Mi|Gi)?B)', re.MULTILINE
|
|
)
|
|
|
|
|
|
def convert_size(size_str: str) -> int:
|
|
base_size = float(size_str.split(' ')[0])
|
|
|
|
if size_str.endswith('GiB'):
|
|
scale = 8589934592
|
|
elif size_str.endswith('MiB'):
|
|
scale = 8388608
|
|
else:
|
|
scale = 1
|
|
return int(round(base_size * scale))
|
|
|
|
|
|
def field_gen(result_str: str):
|
|
for m in field_regex.finditer(result_str):
|
|
d = m.groupdict()
|
|
d['blobs'] = int(d['blobs'])
|
|
d['size'] = convert_size(d['size'])
|
|
yield d
|
|
|
|
|
|
def prune(loki_url: str = None, dry_run: bool = False):
|
|
cmd = ['restic', 'prune']
|
|
if dry_run:
|
|
cmd.append('--dry-run')
|
|
logger.debug(f'Running cmd [bright_black]{" ".join(cmd)}[/]')
|
|
|
|
with console.status('Pruning...'):
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
|
|
if result.returncode != 0:
|
|
logger.error(result.stderr)
|
|
sys.exit(1)
|
|
|
|
d = {f['name']: {'blobs': f['blobs'], 'size': f['size']} for f in field_gen(result.stdout)}
|
|
logger.debug(json.dumps(d, indent=4))
|
|
|
|
if loki_url is not None and not dry_run:
|
|
send_to_loki(loki_url, line=json.dumps(d), backup='prune')
|
|
|
|
|
|
@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)]
|
|
)
|
|
|
|
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__':
|
|
main()
|