Merge pull request #6 from dnephin/exclude_some_images

Support excluding some images from removal

README.rst

@@ -24,6 +24,22 @@ Example:
     dcgc --max-container-age 3days --max-image-age 30days
 
 
+Prevent images from being removed
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+``dcgc`` supports an image exclude list. If you have images that you'd like
+to keep around forever you can use the exclude list to prevent them from
+being removed.
+
+    --exclude-image
+        Never remove images with this tag. May be specified more than once.
+
+    --exclude-image-file
+        Path to a file which contains a list of images to exclude, one
+        image tag per line.
+
+
+
 dcstop
 ------
 

docker_custodian/__about__.py

@@ -1,4 +1,4 @@
 # -*- coding: utf8 -*-
 
-__version_info__ = (0, 4, 0)
+__version_info__ = (0, 5, 0)
 __version__ = '%d.%d.%d' % __version_info__

docker_custodian/docker_gc.py

@@ -68,17 +68,28 @@ def get_all_images(client):
     return images
 
 
-def cleanup_images(client, max_image_age, dry_run):
+def cleanup_images(client, max_image_age, dry_run, exclude_set):
     # re-fetch container list so that we don't include removed containers
     image_tags_in_use = set(
         container['Image'] for container in get_all_containers(client))
 
     images = filter_images_in_use(get_all_images(client), image_tags_in_use)
+    images = filter_excluded_images(images, exclude_set)
 
-    for image_summary in reversed(images):
+    for image_summary in reversed(list(images)):
         remove_image(client, image_summary, max_image_age, dry_run)
 
 
+def filter_excluded_images(images, exclude_set):
+    def include_image(image_summary):
+        image_tags = image_summary.get('RepoTags')
+        if no_image_tags(image_tags):
+            return True
+        return not set(image_tags) & exclude_set
+
+    return filter(include_image, images)
+
+
 def filter_images_in_use(images, image_tags_in_use):
     def get_tag_set(image_summary):
         image_tags = image_summary.get('RepoTags')
@@ -87,10 +98,10 @@ def filter_images_in_use(images, image_tags_in_use):
             return set(['%s:latest' % image_summary['Id'][:12]])
         return set(image_tags)
 
-    def image_is_in_use(image_summary):
+    def image_not_in_use(image_summary):
         return not get_tag_set(image_summary) & image_tags_in_use
 
-    return list(filter(image_is_in_use, images))
+    return filter(image_not_in_use, images)
 
 
 def is_image_old(image, min_date):
@@ -140,22 +151,37 @@ def format_image(image, image_summary):
     return "%s %s" % (image['Id'][:16], get_tags())
 
 
+def build_exclude_set(image_tags, exclude_file):
+    exclude_set = set(image_tags or [])
+
+    def is_image_tag(line):
+        return line and not line.startswith('#')
+
+    if exclude_file:
+        lines = [line.strip() for line in exclude_file.read().split('\n')]
+        exclude_set.update(filter(is_image_tag, lines))
+    return exclude_set
+
+
 def main():
     logging.basicConfig(
         level=logging.INFO,
         format="%(message)s",
         stream=sys.stdout)
 
-    opts = get_opts()
-    client = docker.Client(timeout=opts.timeout)
+    args = get_args()
+    client = docker.Client(timeout=args.timeout)
 
-    if opts.max_container_age:
-        cleanup_containers(client, opts.max_container_age, opts.dry_run)
-    if opts.max_image_age:
-        cleanup_images(client, opts.max_image_age, opts.dry_run)
+    if args.max_container_age:
+        cleanup_containers(client, args.max_container_age, args.dry_run)
+    if args.max_image_age:
+        exclude_set = build_exclude_set(
+            args.exclude_image,
+            args.exclude_image_file)
+        cleanup_images(client, args.max_image_age, args.dry_run, exclude_set)
 
 
-def get_opts(args=None):
+def get_args(args=None):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         '--max-container-age',
@@ -175,6 +201,16 @@ def get_opts(args=None):
     parser.add_argument(
         '-t', '--timeout', type=int, default=60,
         help="HTTP timeout in seconds for making docker API calls.")
+    parser.add_argument(
+        '--exclude-image',
+        action='append',
+        help="Never remove images with this tag.")
+    parser.add_argument(
+        '--exclude-image-file',
+        type=argparse.FileType('r'),
+        help="Path to a file which contains a list of images to exclude, one "
+             "image tag per line.")
+
     return parser.parse_args(args=args)
 
 

tests/docker_gc_test.py

@@ -1,3 +1,6 @@
+from six import StringIO
+import textwrap
+
 import docker.errors
 try:
     from unittest import mock
@@ -63,7 +66,7 @@ def test_cleanup_images(mock_client, now):
     ]
     mock_client.inspect_image.side_effect = iter(mock_images)
 
-    docker_gc.cleanup_images(mock_client, max_image_age, False)
+    docker_gc.cleanup_images(mock_client, max_image_age, False, set())
     assert mock_client.remove_image.mock_calls == [
         mock.call(image['Id']) for image in reversed(images)
     ]
@@ -90,7 +93,29 @@ def test_filter_images_in_use():
         dict(RepoTags=['new_image:latest', 'new_image:123']),
     ]
     actual = docker_gc.filter_images_in_use(images, image_tags_in_use)
-    assert actual == expected
+    assert list(actual) == expected
+
+
+def test_filter_excluded_images():
+    exclude_set = set([
+        'user/one:latest',
+        'user/foo:latest',
+        'other:12345',
+    ])
+    images = [
+        dict(RepoTags=['<none>:<none>'], Id='babababababaabababab'),
+        dict(RepoTags=['user/one:latest', 'user/one:abcd']),
+        dict(RepoTags=['other:abcda']),
+        dict(RepoTags=['other:12345']),
+        dict(RepoTags=['new_image:latest', 'new_image:123']),
+    ]
+    expected = [
+        dict(RepoTags=['<none>:<none>'], Id='babababababaabababab'),
+        dict(RepoTags=['other:abcda']),
+        dict(RepoTags=['new_image:latest', 'new_image:123']),
+    ]
+    actual = docker_gc.filter_excluded_images(images, exclude_set)
+    assert list(actual) == expected
 
 
 def test_is_image_old(image, now):
@@ -178,20 +203,20 @@ def days_as_seconds(num):
     return num * 60 * 60 * 24
 
 
-def test_get_opts_with_defaults():
-    opts = docker_gc.get_opts(args=[])
+def test_get_args_with_defaults():
+    opts = docker_gc.get_args(args=[])
     assert opts.timeout == 60
     assert opts.dry_run is False
     assert opts.max_container_age is None
     assert opts.max_image_age is None
 
 
-def test_get_opts_with_args():
+def test_get_args_with_args():
     with mock.patch(
         'docker_custodian.docker_gc.timedelta_type',
         autospec=True
     ) as mock_timedelta_type:
-        opts = docker_gc.get_opts(args=[
+        opts = docker_gc.get_args(args=[
             '--max-image-age', '30 days',
             '--max-container-age', '3d',
         ])
@@ -224,16 +249,46 @@ def test_get_all_images(mock_client):
     mock_log.info.assert_called_with("Found %s images", count)
 
 
+def test_build_exclude_set():
+    image_tags = [
+        'some_image:latest',
+        'repo/foo:12345',
+        'duplicate:latest',
+    ]
+    exclude_image_file = StringIO(textwrap.dedent("""
+        # Exclude this one because
+        duplicate:latest
+        # Also this one
+        repo/bar:abab
+    """))
+    expected = set([
+        'some_image:latest',
+        'repo/foo:12345',
+        'duplicate:latest',
+        'repo/bar:abab',
+    ])
+
+    exclude_set = docker_gc.build_exclude_set(image_tags, exclude_image_file)
+    assert exclude_set == expected
+
+
+def test_build_exclude_set_empty():
+    exclude_set = docker_gc.build_exclude_set(None, None)
+    assert exclude_set == set()
+
+
 def test_main(mock_client):
     with mock.patch(
             'docker_custodian.docker_gc.docker.Client',
             return_value=mock_client):
 
         with mock.patch(
-                'docker_custodian.docker_gc.get_opts',
-                autospec=True) as mock_get_opts:
-            mock_get_opts.return_value = mock.Mock(
+                'docker_custodian.docker_gc.get_args',
+                autospec=True) as mock_get_args:
+            mock_get_args.return_value = mock.Mock(
                 max_image_age=100,
                 max_container_age=200,
+                exclude_image=[],
+                exclude_image_file=None,
             )
             docker_gc.main()