From 89c9f15bdec9a8a8dcc792f3fe3495c730d8357b Mon Sep 17 00:00:00 2001
From: Priya Wadhwa <priyawadhwa@google.com>
Date: Wed, 13 Jun 2018 10:39:52 -0700
Subject: [PATCH] Add --single-snapshot flag to snapshot once after the build

---
 cmd/executor/cmd/root.go        | 10 +++++++++-
 integration/integration_test.go | 10 +++++++++-
 pkg/executor/executor.go        | 29 +++++++++++++++++++++--------
 3 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/cmd/executor/cmd/root.go b/cmd/executor/cmd/root.go
index 43dbc547c..7de1ead9f 100644
--- a/cmd/executor/cmd/root.go
+++ b/cmd/executor/cmd/root.go
@@ -40,6 +40,7 @@ var (
 	force                       bool
 	buildArgs                   multiArg
 	tarPath                     string
+	singleSnapshot              bool
 )
 
 func init() {
@@ -54,6 +55,7 @@ func init() {
 	RootCmd.PersistentFlags().StringVarP(&logLevel, "verbosity", "v", constants.DefaultLogLevel, "Log level (debug, info, warn, error, fatal, panic")
 	RootCmd.PersistentFlags().BoolVarP(&force, "force", "", false, "Force building outside of a container")
 	RootCmd.PersistentFlags().StringVarP(&tarPath, "tarPath", "", "", "Path to save the image in as a tarball instead of pushing")
+	RootCmd.PersistentFlags().BoolVarP(&singleSnapshot, "single-snapshot", "", false, "Set this flag to take a single snapshot at the end of the build.")
 }
 
 var RootCmd = &cobra.Command{
@@ -79,7 +81,13 @@ var RootCmd = &cobra.Command{
 			logrus.Error(err)
 			os.Exit(1)
 		}
-		ref, image, err := executor.DoBuild(dockerfilePath, srcContext, snapshotMode, buildArgs)
+		ref, image, err := executor.DoBuild(executor.KanikoBuildArgs{
+			DockerfilePath: dockerfilePath,
+			SrcContext:     srcContext,
+			SnapshotMode:   snapshotMode,
+			Args:           buildArgs,
+			SingleSnapshot: singleSnapshot,
+		})
 		if err != nil {
 			logrus.Error(err)
 			os.Exit(1)
diff --git a/integration/integration_test.go b/integration/integration_test.go
index 06e40e07f..238152708 100644
--- a/integration/integration_test.go
+++ b/integration/integration_test.go
@@ -106,6 +106,12 @@ func TestRun(t *testing.T) {
 		"Dockerfile_test_multistage": {"file=/foo2"},
 	}
 
+	// Map for additional flags
+	additionalFlagsMap := map[string][]string{
+		"Dockerfile_test_add":     {"--single-snapshot"},
+		"Dockerfile_test_scratch": {"--single-snapshot"},
+	}
+
 	bucketContextTests := []string{"Dockerfile_test_copy_bucket"}
 
 	// TODO: remove test_user_run from this when https://github.com/GoogleContainerTools/container-diff/issues/237 is fixed
@@ -130,6 +136,7 @@ func TestRun(t *testing.T) {
 				buildArgs = append(buildArgs, buildArgFlag)
 				buildArgs = append(buildArgs, arg)
 			}
+
 			// build docker image
 			dockerImage := strings.ToLower(testRepo + dockerPrefix + dockerfile)
 			dockerCmd := exec.Command("docker",
@@ -152,6 +159,7 @@ func TestRun(t *testing.T) {
 			}
 
 			// build kaniko image
+			additionalFlags := append(buildArgs, additionalFlagsMap[dockerfile]...)
 			kanikoImage := strings.ToLower(testRepo + kanikoPrefix + dockerfile)
 			kanikoCmd := exec.Command("docker",
 				append([]string{"run",
@@ -161,7 +169,7 @@ func TestRun(t *testing.T) {
 					"-f", path.Join(buildContextPath, dockerfilesPath, dockerfile),
 					"-d", kanikoImage,
 					contextFlag, contextPath},
-					buildArgs...)...,
+					additionalFlags...)...,
 			)
 
 			RunCommand(kanikoCmd, t)
diff --git a/pkg/executor/executor.go b/pkg/executor/executor.go
index d5f80b060..85dd7a9d3 100644
--- a/pkg/executor/executor.go
+++ b/pkg/executor/executor.go
@@ -45,9 +45,18 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-func DoBuild(dockerfilePath, srcContext, snapshotMode string, args []string) (name.Reference, v1.Image, error) {
+// KanikoBuildArgs contains all the args required to build the image
+type KanikoBuildArgs struct {
+	DockerfilePath string
+	SrcContext     string
+	SnapshotMode   string
+	Args           []string
+	SingleSnapshot bool
+}
+
+func DoBuild(k KanikoBuildArgs) (name.Reference, v1.Image, error) {
 	// Parse dockerfile and unpack base image to root
-	d, err := ioutil.ReadFile(dockerfilePath)
+	d, err := ioutil.ReadFile(k.DockerfilePath)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -58,12 +67,12 @@ func DoBuild(dockerfilePath, srcContext, snapshotMode string, args []string) (na
 	}
 	dockerfile.ResolveStages(stages)
 
-	hasher, err := getHasher(snapshotMode)
+	hasher, err := getHasher(k.SnapshotMode)
 	if err != nil {
 		return nil, nil, err
 	}
 	for index, stage := range stages {
-		baseImage, err := util.ResolveEnvironmentReplacement(stage.BaseName, args, false)
+		baseImage, err := util.ResolveEnvironmentReplacement(stage.BaseName, k.Args, false)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -109,9 +118,10 @@ func DoBuild(dockerfilePath, srcContext, snapshotMode string, args []string) (na
 		if err := resolveOnBuild(&stage, &imageConfig.Config); err != nil {
 			return nil, nil, err
 		}
-		buildArgs := dockerfile.NewBuildArgs(args)
-		for _, cmd := range stage.Commands {
-			dockerCommand, err := commands.GetCommand(cmd, srcContext)
+		buildArgs := dockerfile.NewBuildArgs(k.Args)
+		for index, cmd := range stage.Commands {
+			finalCmd := index == len(stage.Commands)-1
+			dockerCommand, err := commands.GetCommand(cmd, k.SrcContext)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -121,11 +131,14 @@ func DoBuild(dockerfilePath, srcContext, snapshotMode string, args []string) (na
 			if err := dockerCommand.ExecuteCommand(&imageConfig.Config, buildArgs); err != nil {
 				return nil, nil, err
 			}
-			if !finalStage {
+			if !finalStage || (k.SingleSnapshot && !finalCmd) {
 				continue
 			}
 			// Now, we get the files to snapshot from this command and take the snapshot
 			snapshotFiles := dockerCommand.FilesToSnapshot()
+			if k.SingleSnapshot && finalCmd {
+				snapshotFiles = nil
+			}
 			contents, err := snapshotter.TakeSnapshot(snapshotFiles)
 			if err != nil {
 				return nil, nil, err