bundles/zfs: refactor zfs-auto-snapshot

bundles/zfs/files/zfs-auto-snapshot

@@ -9,21 +9,42 @@ from subprocess import check_call, check_output
 from sys import argv
 
 
-def create_snap_and_rotate(ds, label, retain, now, all_snapshots):
-    new_snap = '{}@zfs-auto-snap_{}-{}'.format(ds, label, now)
-    check_call(['zfs', 'snapshot', new_snap])
+def list_datasets():
+    datasets = set()
+    for line in check_output(['zfs', 'list', '-H', '-o', 'name']).splitlines():
+        line = line.decode('UTF-8')
+
+        for prefix in metadata.get('snapshot_never', set()):
+            if line.startswith(prefix):
+                break
+        else:
+            datasets.add(line)
+
+    return datasets
+
+
+def get_filtered_snapshots_for_dataset(ds):
+    all_snapshots = check_output(['zfs', 'list', '-H', '-o', 'name', '-t', 'snapshot', ds]).decode('UTF-8').splitlines()
 
     prefix = '{}@zfs-auto-snap_{}-'.format(ds, label)
-    my_candidates = []
+    snapshots = set()
+
     for i in sorted(all_snapshots):
         if i.startswith(prefix):
-            my_candidates.append(i)
+            snapshots.add(i)
 
-    my_candidates.append(new_snap)
+    return snapshots
 
-    for i in my_candidates[:-retain]:
-        assert '@' in i, 'BUG! Dataset "{}" has no @!'.format(i)
-        check_call(['zfs', 'destroy', i])
+
+def delete_snapshot(snap):
+    assert '@' in snap, 'BUG! Dataset "{}" has no @!'.format(snap)
+    print('deleting snapshot {}'.format(snap))
+    check_call(['zfs', 'destroy', snap])
+
+
+def create_snapshot(ds, label, now):
+    check_call(['zfs', 'snapshot', '{}@zfs-auto-snap_{}-{}'.format(ds, label, now)])
+    print('created snapshot {}@zfs-auto-snap_{}-{}'.format(ds, label, now))
 
 
 label = argv[1]
@@ -31,28 +52,33 @@ label = argv[1]
 with open('/etc/zfs-snapshot-config.json', 'r') as fp:
     metadata = loads(fp.read())
 
-datasets = set()
-for line in check_output(['zfs', 'list', '-H', '-o', 'name']).splitlines():
-    line = line.decode('UTF-8')
-
-    for prefix in metadata.get('snapshot_never', set()):
-        if line.startswith(prefix):
-            break
-    else:
-        datasets.add(line)
-
 default_retain = metadata['retain_defaults'][label]
 now = datetime.now().strftime('%F-%H%M')
 snapshots_created = False
 
-if datasets:
-    all_snapshots = check_output(['zfs', 'list', '-H', '-o', 'name', '-t', 'snap']).decode('UTF-8').splitlines()
+for ds in list_datasets():
+    retain = int(metadata.get('retain_per_dataset', {}).get(ds, {}).get(label, default_retain))
 
-    for ds in datasets:
-        retain = int(metadata.get('retain_per_dataset', {}).get(ds, {}).get(label, default_retain))
-        if retain > 0:
-            create_snap_and_rotate(ds, label, retain, now, all_snapshots)
-            snapshots_created = True
+    if retain > 0:
+        create_snapshot(ds, label, now)
+        snapshots_created = True
+
+    existing_snapshots = get_filtered_snapshots_for_dataset(ds)
+
+    if retain > 0:
+        # Why +1 here? Because we're specifying the amount of hours we want
+        # to go back in time, not the amount of snapshots we want to keep.
+        # Stating '1 month' does mean 'i want to be able to go back atleast
+        # one monthly snapshot', not 'keep one monthly snapshot'. If we only
+        # kept one snapshot, that wouldn't be possible for most of the time,
+        # because the monthly snapshot is actually less than a month old.
+        snapshots_to_keep = retain+1
+
+        for snapshot in sorted(existing_snapshots)[:-snapshots_to_keep]:
+            delete_snapshot(snapshot)
+    else:
+        for snapshot in existing_snapshots:
+            delete_snapshot(snapshot)
 
 with open('/var/tmp/zfs-auto-snapshot.status', 'w') as fp:
     fp.write('{}\n'.format(datetime.now().strftime('%s') if snapshots_created else 0))