From 4a3941041b328a292d586fa22f6b42907d47b771 Mon Sep 17 00:00:00 2001
From: Anton Lundin <glance@acc.umu.se>
Date: Mon, 28 Oct 2019 11:04:51 +0100
Subject: [PATCH] Implement volume exclusion

This implements --exclude-volume and --exclude-volume-file, so one can
list volumes that shouldn't be removed.

Signed-off-by: Anton Lundin <glance@acc.umu.se>
---
 docker_custodian/docker_gc.py | 31 ++++++++++++++++--
 tests/docker_gc_test.py       | 60 +++++++++++++++++++++++++++++++++--
 2 files changed, 86 insertions(+), 5 deletions(-)

diff --git a/docker_custodian/docker_gc.py b/docker_custodian/docker_gc.py
index dffcd03..f958cac 100644
--- a/docker_custodian/docker_gc.py
+++ b/docker_custodian/docker_gc.py
@@ -228,10 +228,23 @@ def remove_volume(client, volume, dry_run):
     api_call(client.remove_volume, name=volume['Name'])
 
 
-def cleanup_volumes(client, dry_run):
+def filter_excluded_volumes(volumes, exclude_set):
+    def include_volume(volume):
+        image_name = volume['Name']
+        for exclude_pattern in exclude_set:
+            if fnmatch.fnmatch(image_name, exclude_pattern):
+                return False
+        return True
+
+    return filter(include_volume, volumes)
+
+
+def cleanup_volumes(client, dry_run, exclude_set):
     dangling_volumes = get_dangling_volumes(client)
 
-    for volume in reversed(dangling_volumes):
+    dangling_volumes = filter_excluded_volumes(dangling_volumes, exclude_set)
+
+    for volume in dangling_volumes:
         log.info("Removing dangling volume %s", volume['Name'])
         remove_volume(client, volume, dry_run)
 
@@ -317,7 +330,10 @@ def main():
         cleanup_images(client, args.max_image_age, args.dry_run, exclude_set)
 
     if args.dangling_volumes:
-        cleanup_volumes(client, args.dry_run)
+        exclude_set = build_exclude_set(
+            args.exclude_volume,
+            args.exclude_volume_file)
+        cleanup_volumes(client, args.dry_run, exclude_set)
 
 
 def get_args(args=None):
@@ -353,6 +369,15 @@ def get_args(args=None):
         type=argparse.FileType('r'),
         help="Path to a file which contains a list of images to exclude, one "
              "image tag per line.")
+    parser.add_argument(
+        '--exclude-volume',
+        action='append',
+        help="Never remove volume with this name.")
+    parser.add_argument(
+        '--exclude-volume-file',
+        type=argparse.FileType('r'),
+        help="Path to a file which contains a list of volumes to exclude, one "
+             "image tag per line.")
     parser.add_argument(
         '--exclude-container-label',
         action='append', type=str, default=[],
diff --git a/tests/docker_gc_test.py b/tests/docker_gc_test.py
index 445301f..90a7deb 100644
--- a/tests/docker_gc_test.py
+++ b/tests/docker_gc_test.py
@@ -149,13 +149,67 @@ def test_cleanup_volumes(mock_client):
         'Warnings': None,
     }
 
-    docker_gc.cleanup_volumes(mock_client, False)
+    docker_gc.cleanup_volumes(mock_client, False, set())
     assert mock_client.remove_volume.mock_calls == [
         mock.call(name=volume['Name'])
-        for volume in reversed(volumes['Volumes'])
+        for volume in volumes['Volumes']
     ]
 
 
+def test_filter_cleanup_volumes(mock_client):
+    mock_client.volumes.return_value = {
+        'Volumes': [
+            {
+                'Mountpoint': 'unused',
+                'Labels': None,
+                'Driver': 'unused',
+                'Name': u'unused'
+            },
+            {
+                'Mountpoint': 'filtered',
+                'Labels': None,
+                'Driver': 'unused',
+                'Name': u'filtered'
+            },
+        ],
+        'Warnings': None,
+    }
+
+    docker_gc.cleanup_volumes(mock_client, False, set(['filtered']))
+    assert mock_client.remove_volume.mock_calls == [
+        mock.call(name='unused')
+    ]
+
+
+def test_filter_excluded_volumes(mock_client):
+    exclude_set = set(['filtered'])
+    volumes = [
+        {
+            'Mountpoint': 'unused',
+            'Labels': None,
+            'Driver': 'unused',
+            'Name': u'unused'
+        },
+        {
+            'Mountpoint': 'filtered',
+            'Labels': None,
+            'Driver': 'unused',
+            'Name': u'filtered'
+        },
+    ]
+    expected = [
+        {
+            'Mountpoint': 'unused',
+            'Labels': None,
+            'Driver': 'unused',
+            'Name': u'unused'
+        },
+    ]
+
+    actual = docker_gc.filter_excluded_volumes(volumes, exclude_set)
+    assert list(actual) == expected
+
+
 def test_filter_images_in_use():
     image_tags_in_use = set([
         'user/one:latest',
@@ -527,6 +581,8 @@ def test_main(mock_client):
                 max_container_age=200,
                 exclude_image=[],
                 exclude_image_file=None,
+                exclude_volume=[],
+                exclude_volume_file=None,
                 exclude_container_label=[],
             )
             docker_gc.main()