Repository: goldmann/docker-squash

Issue:
Remove squashed layers from Docker daemon
In most cases we're not interested in layers that were squashed. In case where the input and target image has the same tag - all squashed layers will be "orphaned" leaving them in Docker daemon and wasting space. We should probably include a `--cleanup` switch or something similar to remove layers without tag that were "orphaned" during the squashing process.

Discussion:

Patch:
diff --git a/docker_squash/cli.py b/docker_squash/cli.py
index 462b872..cff0773 100644
--- a/docker_squash/cli.py
+++ b/docker_squash/cli.py
@@ -42,7 +42,9 @@ class CLI(object):
         parser.add_argument(
             '-f', '--from-layer', help='ID of the layer or image ID or image name. If not specified will squash all layers in the image')
         parser.add_argument(
-            '-t', '--tag', help="Specify the tag to be used for the new image. By default it'll be set to 'image' argument")
+            '-t', '--tag', help="Specify the tag to be used for the new image. by default it'll be set to 'image' argument")
+        parser.add_argument(
+            '-c', '--cleanup', action='store_true', help="Remove source image from Docker after squashing")
         parser.add_argument(
             '--tmp-dir', help='Temporary directory to be used')
         parser.add_argument(
@@ -59,7 +61,7 @@ class CLI(object):
 
         try:
             squash.Squash(log=self.log, image=args.image,
-                          from_layer=args.from_layer, tag=args.tag, output_path=args.output_path, tmp_dir=args.tmp_dir, development=args.development).run()
+                          from_layer=args.from_layer, tag=args.tag, output_path=args.output_path, tmp_dir=args.tmp_dir, development=args.development, cleanup=args.cleanup).run()
         except KeyboardInterrupt:
             self.log.error("Program interrupted by user, exiting...")
             sys.exit(1)
diff --git a/docker_squash/squash.py b/docker_squash/squash.py
index d6d21e0..e81960a 100644
--- a/docker_squash/squash.py
+++ b/docker_squash/squash.py
@@ -14,7 +14,7 @@ from docker_squash.version import version
 class Squash(object):
 
     def __init__(self, log, image, docker=None, from_layer=None, tag=None, tmp_dir=None,
-                 output_path=None, load_image=True, development=False):
+                 output_path=None, load_image=True, development=False, cleanup=False):
         self.log = log
         self.docker = docker
         self.image = image
@@ -24,6 +24,7 @@ class Squash(object):
         self.output_path = output_path
         self.load_image = load_image
         self.development = development
+        self.cleanup = cleanup
 
         if not docker:
             self.docker = common.docker_client(self.log)
@@ -65,12 +66,22 @@ class Squash(object):
 
             raise
 
+    def _cleanup(self):
+        try:
+            image_id = self.docker.inspect_image(self.image)['Id']
+        except:
+            self.log.warn("Could not get the image ID for %s image, skipping cleanup after squashing" % self.image)
+            return
+
+        self.log.info("Removing old %s image..." % self.image)
+        self.docker.remove_image(image_id, force=False, noprune=False)
+        self.log.info("Image removed!")
+
     def squash(self, image):
         # Do the actual squashing
         new_image_id = image.squash()
 
-        self.log.info(
-            "New squashed image ID is %s" % new_image_id)
+        self.log.info("New squashed image ID is %s" % new_image_id)
 
         if self.output_path:
             # Move the tar archive to the specified path
@@ -80,8 +91,16 @@ class Squash(object):
             # Load squashed image into Docker
             image.load_squashed_image()
 
+        # Clean up all temporary files
         image.cleanup()
 
+        # Remove the source image - this is the only possible way
+        # to remove orphaned layers from Docker daemon at the build time.
+        # We cannot use here a tag name because it could be used as the target,
+        # squashed image tag - we need to use the image ID.
+        if self.cleanup:
+            self._cleanup()
+
         self.log.info("Done")
 
         return new_image_id
diff --git a/docker_squash/v2_image.py b/docker_squash/v2_image.py
index c2fa1ac..4df00d7 100644
--- a/docker_squash/v2_image.py
+++ b/docker_squash/v2_image.py
@@ -280,10 +280,10 @@ class V2Image(Image):
         if self.layer_paths_to_move:
             config['parent'] = self.layer_paths_to_move[-1]
         else:
-            del config['parent']
+            config.pop("parent", None)
         # Update 'id' - it should be the path to the layer
         config['id'] = layer_path_id
-        del config['container']
+        config.pop("container", None)
         return config
 
     def _generate_image_metadata(self):
@@ -294,7 +294,7 @@ class V2Image(Image):
         metadata['created'] = self.date
 
         # Remove unnecessary or old fields
-        del metadata['container']
+        metadata.pop("container", None)
 
         # Remove squashed layers from history
         metadata['history'] = metadata['history'][:len(self.layers_to_move)]


Test patch:
diff --git a/tests/test_unit_squash.py b/tests/test_unit_squash.py
new file mode 100644
index 0000000..c136761
--- /dev/null
+++ b/tests/test_unit_squash.py
@@ -0,0 +1,43 @@
+import unittest
+import mock
+import six
+import logging
+
+from docker_squash.squash import Squash
+from docker_squash.image import Image
+from docker_squash.errors import SquashError
+
+
+class TestSquash(unittest.TestCase):
+    
+    def setUp(self):
+        self.log = mock.Mock()
+        self.docker_client = mock.Mock()
+        self.docker_client.version.return_value = {'GitCommit': "commit/9.9.9", 'ApiVersion': "9.99"}
+
+    def test_handle_case_when_no_image_is_provided(self):
+        squash = Squash(self.log, None, self.docker_client)
+        with self.assertRaises(SquashError) as cm:
+            squash.run()
+        self.assertEquals(
+            str(cm.exception), "Image is not provided")
+
+    def test_exit_if_no_output_path_provided_and_loading_is_disabled_too(self):
+        squash = Squash(self.log, 'image', self.docker_client, load_image=False, output_path=None)
+        squash.run()
+        self.log.warn.assert_called_with("No output path specified and loading into Docker is not selected either; squashed image would not accessible, proceeding with squashing doesn't make sense")
+
+    @mock.patch('docker_squash.squash.V2Image')
+    def test_should_not_cleanup_after_squashing(self, v2_image):
+        squash = Squash(self.log, 'image', self.docker_client, load_image=True)
+        squash.run()
+    
+        v2_image.cleanup.assert_not_called()
+
+    @mock.patch('docker_squash.squash.V2Image')
+    def test_should_cleanup_after_squashing(self, v2_image):
+        self.docker_client.inspect_image.return_value = {'Id': "abcdefgh"}
+        squash = Squash(self.log, 'image', self.docker_client, load_image=True, cleanup=True)
+        squash.run()
+    
+        self.docker_client.remove_image.assert_called_with('abcdefgh', force=False, noprune=False)
diff --git a/tests/test_unit_v1_image.py b/tests/test_unit_v1_image.py
index 6316399..07081ab 100644
--- a/tests/test_unit_v1_image.py
+++ b/tests/test_unit_v1_image.py
@@ -299,24 +299,5 @@ class TestAddMarkers(unittest.TestCase):
         self.assertTrue(tar_info.isfile())
 
 
-class TestGeneral(unittest.TestCase):
-
-    def setUp(self):
-        self.log = mock.Mock()
-        self.docker_client = mock.Mock()
-        self.docker_client.version.return_value = {'GitCommit': "commit/9.9.9", 'ApiVersion': "9.99"}
-
-    def test_handle_case_when_no_image_is_provided(self):
-        squash = Squash(self.log, None, self.docker_client)
-        with self.assertRaises(SquashError) as cm:
-            squash.run()
-        self.assertEquals(
-            str(cm.exception), "Image is not provided")
-
-    def test_exit_if_no_output_path_provided_and_loading_is_disabled_too(self):
-        squash = Squash(self.log, 'image', self.docker_client, load_image=False, output_path=None)
-        squash.run()
-        self.log.warn.assert_called_with("No output path specified and loading into Docker is not selected either; squashed image would not accessible, proceeding with squashing doesn't make sense")
-
 if __name__ == '__main__':
     unittest.main()