From 7a23303a1bea814fca8708672c8f9b1c9e99de99 Mon Sep 17 00:00:00 2001
From: jay <jay@lawsnote.com>
Date: Tue, 19 Jul 2022 16:42:05 +0800
Subject: [PATCH] [feat] add rm option

---
 cmd/main.go               |  4 +++-
 internal/docker/docker.go | 17 ++++++++++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index e39974a..5679474 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -16,6 +16,7 @@ func main() {
 	var imageName string
 	var tarName string
 	var doPull bool
+	var rm bool
 
 	rootCmd = &cobra.Command{
 		Use:   "docker-save",
@@ -27,7 +28,7 @@ func main() {
 			}
 
 			ctx := context.Background()
-			if err := docker.SaveImageTar(ctx, imageName, outputPath, tarName, doPull); err != nil {
+			if err := docker.SaveImageTar(ctx, imageName, outputPath, tarName, doPull, rm); err != nil {
 				log.Fatal(err)
 			}
 		},
@@ -36,6 +37,7 @@ func main() {
 	rootCmd.Flags().StringVarP(&imageName, "image", "i", "", "image name")
 	rootCmd.Flags().StringVarP(&tarName, "name", "n", "", "tar name")
 	rootCmd.Flags().BoolVarP(&doPull, "pull", "p", false, "pull image")
+	rootCmd.Flags().BoolVarP(&rm, "rm", "", false, "delete image after save")
 	rootCmd.MarkFlagRequired("image")
 	rootCmd.MarkFlagRequired("output")
 
diff --git a/internal/docker/docker.go b/internal/docker/docker.go
index a427159..e866caf 100644
--- a/internal/docker/docker.go
+++ b/internal/docker/docker.go
@@ -12,7 +12,7 @@ import (
 	"github.com/docker/docker/client"
 )
 
-func SaveImageTar(ctx context.Context, image, path, name string, doPull bool) error {
+func SaveImageTar(ctx context.Context, image, path, name string, doPull, rm bool) error {
 
 	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
 	if err != nil {
@@ -30,6 +30,12 @@ func SaveImageTar(ctx context.Context, image, path, name string, doPull bool) er
 		return fmt.Errorf("save image %s failed: %v", image, err)
 	}
 
+	if rm {
+		if err := rmImage(ctx, cli, image); err != nil {
+			return fmt.Errorf("remove image %s failed: %v", image, err)
+		}
+	}
+
 	return nil
 }
 
@@ -46,6 +52,15 @@ func pullImage(ctx context.Context, cli *client.Client, image string) error {
 	return nil
 }
 
+func rmImage(ctx context.Context, cli *client.Client, image string) error {
+	_, err := cli.ImageRemove(ctx, image, types.ImageRemoveOptions{})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func saveToTar(ctx context.Context, cli *client.Client, image, path, name string) error {
 	out, err := cli.ImageSave(ctx, []string{image})
 	if err != nil {