tk export: Expand merging capabilities with --merge-strategy flag

cmd/tk/export.go

@@ -34,19 +34,31 @@ func exportCmd() *cli.Command {
 	)
 
 	extension := cmd.Flags().String("extension", "yaml", "File extension")
-	merge := cmd.Flags().Bool("merge", false, "Allow merging with existing directory")
 	parallel := cmd.Flags().IntP("parallel", "p", 8, "Number of environments to process in parallel")
 	cachePath := cmd.Flags().StringP("cache-path", "c", "", "Local file path where cached evaluations should be stored")
 	cacheEnvs := cmd.Flags().StringArrayP("cache-envs", "e", nil, "Regexes which define which environment should be cached (if caching is enabled)")
 	ballastBytes := cmd.Flags().Int("mem-ballast-size-bytes", 0, "Size of memory ballast to allocate. This may improve performance for large environments.")
 
+	merge := cmd.Flags().Bool("merge", false, "Allow merging with existing directory")
+	if err := cmd.Flags().MarkDeprecated("merge", "use --merge-strategy=fail-on-conflicts instead"); err != nil {
+		panic(err)
+	}
+	mergeStrategy := cmd.Flags().String("merge-strategy", "", "What to do when exporting to an existing directory. Values: 'fail-on-conficts', 'replace-envs'")
+
 	vars := workflowFlags(cmd.Flags())
 	getJsonnetOpts := jsonnetFlags(cmd.Flags())
 	getLabelSelector := labelSelectorFlag(cmd.Flags())
 
 	recursive := cmd.Flags().BoolP("recursive", "r", false, "Look recursively for Tanka environments")
 
 	cmd.Run = func(cmd *cli.Command, args []string) error {
+		// `--merge` is deprecated in favor of `--merge-strategy`. However, merge has to keep working for now.
+		if *merge && *mergeStrategy != "" {
+			panic("cannot use --merge and --merge-strategy at the same time")
+		} else if *merge {
+			*mergeStrategy = string(tanka.ExportMergeStrategyFailConflicts)
+		}
+
 		// Allocate a block of memory to alter GC behaviour. See https://github.com/golang/go/issues/23044
 		ballast := make([]byte, *ballastBytes)
 		defer runtime.KeepAlive(ballast)
@@ -57,9 +69,9 @@ func exportCmd() *cli.Command {
 		}
 
 		opts := tanka.ExportEnvOpts{
-			Format:    *format,
-			Extension: *extension,
-			Merge:     *merge,
+			Format:        *format,
+			Extension:     *extension,
+			MergeStrategy: tanka.ExportMergeStrategy(*mergeStrategy),
 			Opts: tanka.Opts{
 				JsonnetOpts: getJsonnetOpts(),
 				Filters:     filters,

docs/docs/exporting.md

@@ -99,15 +99,29 @@ $ tk export exportDir environments/ --recursive
 $ tk export exportDir environments/ -r -l team=infra
 ```
 
-## Using a memory ballast
+## Performance features
+
+When exporting a large amount of environments, jsonnet evaluation can become a bottleneck. To speed up the process, Tanka provides a few optional features.
+
+### Partial export (in a GitOps context)
+
+Given multiple environments, one may want to only export the environments that were modified since the last export. This is enabled by passing both the `--merge-strategy=replace-envs` flags.
+
+When these flags are passed, Tanka will:
+
+1. Delete the manifests that were previously exported by the environments that are being exported. This is done by looking at the `manifest.json` file that is generated by Tanka when exporting. The related entries are also removed from the `manifest.json` file.
+2. Generate the manifests for the targeted environments into the output directory.
+3. Add in the new manifests entries into the `manifest.json` file and re-export it.
+
+### Using a memory ballast
 
 _Read [this blog post](https://blog.twitch.tv/en/2019/04/10/go-memory-ballast-how-i-learnt-to-stop-worrying-and-love-the-heap/) for more information about memory ballasts._
 
 For large environments that load lots of data into memory on evaluation, a memory ballast can dramatically improve performance. This feature is exposed through the `--mem-ballast-size-bytes` flag on the export command.
 
 Anecdotally (Grafana Labs), environments that took around a minute to load were able to load in around 45 secs with a ballast of 5GB (`--mem-ballast-size-bytes=5368709120`). Decreasing the ballast size resulted in negative impact on performance, and increasing it more did not result in any noticeable impact.
 
-## Caching
+### Caching
 
 Tanka can also cache the results of the export. This is useful if you often export the same files and want to avoid recomputing them. The cache key is calculated from the main file and all of its transitive imports, so any change to any file possibly used in an environment will invalidate the cache.
 

pkg/tanka/export.go

@@ -5,12 +5,15 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"io/fs"
+	"log"
 	"os"
 	"path/filepath"
 	"strings"
 	"text/template"
 
 	"github.com/Masterminds/sprig/v3"
+	"github.com/pkg/errors"
 	"k8s.io/apimachinery/pkg/labels"
 
 	"github.com/grafana/tanka/pkg/kubernetes/manifest"
@@ -27,20 +30,32 @@ const BelRune = string(rune(7))
 // debugging purposes.
 const manifestFile = "manifest.json"
 
+type ExportMergeStrategy string
+
+const (
+	ExportMergeStrategyNone          ExportMergeStrategy = ""
+	ExportMergeStrategyFailConflicts ExportMergeStrategy = "fail-on-conflicts"
+	ExportMergeStrategyReplaceEnvs   ExportMergeStrategy = "replace-envs"
+)
+
 // ExportEnvOpts specify options on how to export environments
 type ExportEnvOpts struct {
 	// formatting the filename based on the exported Kubernetes manifest
 	Format string
 	// extension of the filename
 	Extension string
-	// merge export with existing directory
-	Merge bool
 	// optional: options to parse Jsonnet
 	Opts Opts
 	// optional: filter environments based on labels
 	Selector labels.Selector
 	// optional: number of environments to process in parallel
 	Parallelism int
+
+	// What to do when exporting to an existing directory
+	// - none: fail when directory is not empty
+	// - fail-on-conflicts: fail when an exported file already exists
+	// - replace-envs: delete files previously exported by the targeted envs and re-export them
+	MergeStrategy ExportMergeStrategy
 }
 
 func ExportEnvironments(envs []*v1alpha1.Environment, to string, opts *ExportEnvOpts) error {
@@ -52,8 +67,15 @@ func ExportEnvironments(envs []*v1alpha1.Environment, to string, opts *ExportEnv
 	if err != nil {
 		return fmt.Errorf("checking target dir: %s", err)
 	}
-	if !empty && !opts.Merge {
-		return fmt.Errorf("output dir `%s` not empty. Pass --merge to ignore this", to)
+	if !empty && opts.MergeStrategy == ExportMergeStrategyNone {
+		return fmt.Errorf("output dir `%s` not empty. Pass a different --merge-strategy to ignore this", to)
+	}
+
+	// delete files previously exported by the targeted envs.
+	if opts.MergeStrategy == ExportMergeStrategyReplaceEnvs {
+		if err := deletePreviouslyExportedManifests(to, envs); err != nil {
+			return fmt.Errorf("deleting previously exported manifests: %w", err)
+		}
 	}
 
 	// get all environments for paths
@@ -122,19 +144,7 @@ func ExportEnvironments(envs []*v1alpha1.Environment, to string, opts *ExportEnv
 		}
 	}
 
-	// Write manifest file
-	if len(fileToEnv) != 0 {
-		data, err := json.MarshalIndent(fileToEnv, "", "    ")
-		if err != nil {
-			return err
-		}
-		path := filepath.Join(to, manifestFile)
-		if err := writeExportFile(path, data); err != nil {
-			return err
-		}
-	}
-
-	return nil
+	return exportManifestFile(to, fileToEnv, nil)
 }
 
 func fileExists(name string) (bool, error) {
@@ -164,6 +174,75 @@ func dirEmpty(dir string) (bool, error) {
 	return false, err
 }
 
+func deletePreviouslyExportedManifests(path string, envs []*v1alpha1.Environment) error {
+	fileToEnvMap := make(map[string]string)
+
+	manifestFilePath := filepath.Join(path, manifestFile)
+	manifestContent, err := os.ReadFile(manifestFilePath)
+	if errors.Is(err, fs.ErrNotExist) {
+		log.Printf("Warning: No manifest file found at %s, skipping deletion of previously exported manifests\n", manifestFilePath)
+		return nil
+	} else if err != nil {
+		return err
+	}
+
+	if err := json.Unmarshal(manifestContent, &fileToEnvMap); err != nil {
+		return err
+	}
+
+	envNames := make(map[string]struct{})
+	for _, env := range envs {
+		envNames[env.Metadata.Namespace] = struct{}{}
+	}
+
+	var deletedManifestKeys []string
+	for exportedManifest, manifestEnv := range fileToEnvMap {
+		if _, ok := envNames[manifestEnv]; ok {
+			deletedManifestKeys = append(deletedManifestKeys, exportedManifest)
+			if err := os.Remove(filepath.Join(path, exportedManifest)); err != nil {
+				return err
+			}
+		}
+	}
+
+	return exportManifestFile(path, nil, deletedManifestKeys)
+}
+
+// exportManifestFile writes a manifest file that maps the exported files to their environment.
+// If the file already exists, the new entries will be merged with the existing ones.
+func exportManifestFile(path string, newFileToEnvMap map[string]string, deletedKeys []string) error {
+	if len(newFileToEnvMap) == 0 && len(deletedKeys) == 0 {
+		return nil
+	}
+
+	currentFileToEnvMap := make(map[string]string)
+	manifestFilePath := filepath.Join(path, manifestFile)
+	if manifestContent, err := os.ReadFile(manifestFilePath); err != nil && !errors.Is(err, fs.ErrNotExist) {
+		return fmt.Errorf("reading existing manifest file: %w", err)
+	} else if err == nil {
+		// Only read the manifest file if it exists.
+		// If it doesn't exist, currentFileToEnvMap will be empty, meaning that we're starting from a new export dir.
+		if err := json.Unmarshal(manifestContent, &currentFileToEnvMap); err != nil {
+			return fmt.Errorf("unmarshalling existing manifest file: %w", err)
+		}
+	}
+
+	for k, v := range newFileToEnvMap {
+		currentFileToEnvMap[k] = v
+	}
+	for _, k := range deletedKeys {
+		delete(currentFileToEnvMap, k)
+	}
+
+	// Write manifest file
+	data, err := json.MarshalIndent(currentFileToEnvMap, "", "    ")
+	if err != nil {
+		return fmt.Errorf("marshalling manifest file: %w", err)
+	}
+
+	return writeExportFile(manifestFilePath, data)
+}
+
 func writeExportFile(path string, data []byte) error {
 	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
 		return fmt.Errorf("creating filepath '%s': %s", filepath.Dir(path), err)

pkg/tanka/export_test.go

@@ -1,6 +1,18 @@
 package tanka
 
-import "testing"
+import (
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/grafana/tanka/pkg/jsonnet"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/labels"
+)
 
 func Test_replaceTmplText(t *testing.T) {
 	type args struct {
@@ -31,3 +43,132 @@ func Test_replaceTmplText(t *testing.T) {
 		})
 	}
 }
+
+func TestExportEnvironments(t *testing.T) {
+	tempDir := t.TempDir()
+	require.NoError(t, os.Chdir("testdata"))
+	defer func() { require.NoError(t, os.Chdir("..")) }()
+
+	// Find envs
+	envs, err := FindEnvs("test-export-envs", FindOpts{Selector: labels.Everything()})
+	require.NoError(t, err)
+
+	// Export all envs
+	opts := &ExportEnvOpts{
+		Format:    "{{.metadata.namespace}}/{{.metadata.name}}",
+		Extension: "yaml",
+	}
+	opts.Opts.ExtCode = jsonnet.InjectedCode{
+		"deploymentName": "'initial-deployment'",
+		"serviceName":    "'initial-service'",
+	}
+	require.NoError(t, ExportEnvironments(envs, tempDir, opts))
+	checkFiles(t, tempDir, []string{
+		filepath.Join(tempDir, "inline-namespace1", "my-configmap.yaml"),
+		filepath.Join(tempDir, "inline-namespace1", "my-deployment.yaml"),
+		filepath.Join(tempDir, "inline-namespace1", "my-service.yaml"),
+		filepath.Join(tempDir, "inline-namespace2", "my-deployment.yaml"),
+		filepath.Join(tempDir, "inline-namespace2", "my-service.yaml"),
+		filepath.Join(tempDir, "static", "initial-deployment.yaml"),
+		filepath.Join(tempDir, "static", "initial-service.yaml"),
+		filepath.Join(tempDir, "manifest.json"),
+	})
+	manifestContent, err := os.ReadFile(filepath.Join(tempDir, "manifest.json"))
+	require.NoError(t, err)
+	assert.Equal(t, string(manifestContent), `{
+    "inline-namespace1/my-configmap.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "inline-namespace1/my-deployment.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "inline-namespace1/my-service.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "inline-namespace2/my-deployment.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "inline-namespace2/my-service.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "static/initial-deployment.yaml": "test-export-envs/static-env/main.jsonnet",
+    "static/initial-service.yaml": "test-export-envs/static-env/main.jsonnet"
+}`)
+
+	// Try to re-export
+	assert.EqualError(t, ExportEnvironments(envs, tempDir, opts), fmt.Sprintf("output dir `%s` not empty. Pass a different --merge-strategy to ignore this", tempDir))
+
+	// Try to re-export with the --merge-strategy=fail-on-conflicts flag. Will still fail because Tanka will not overwrite manifests silently
+	opts.MergeStrategy = ExportMergeStrategyFailConflicts
+	assert.ErrorContains(t, ExportEnvironments(envs, tempDir, opts), "already exists. Aborting")
+
+	// Re-export only one env with --merge-stategy=replace-envs flag
+	opts.Opts.ExtCode = jsonnet.InjectedCode{
+		"deploymentName": "'updated-deployment'",
+		"serviceName":    "'updated-service'",
+	}
+	opts.MergeStrategy = ExportMergeStrategyReplaceEnvs
+	staticEnv, err := FindEnvs("test-export-envs", FindOpts{Selector: labels.SelectorFromSet(labels.Set{"type": "static"})})
+	require.NoError(t, err)
+	require.NoError(t, ExportEnvironments(staticEnv, tempDir, opts))
+	checkFiles(t, tempDir, []string{
+		filepath.Join(tempDir, "inline-namespace1", "my-configmap.yaml"),
+		filepath.Join(tempDir, "inline-namespace1", "my-deployment.yaml"),
+		filepath.Join(tempDir, "inline-namespace1", "my-service.yaml"),
+		filepath.Join(tempDir, "inline-namespace2", "my-deployment.yaml"),
+		filepath.Join(tempDir, "inline-namespace2", "my-service.yaml"),
+		filepath.Join(tempDir, "static", "updated-deployment.yaml"),
+		filepath.Join(tempDir, "static", "updated-service.yaml"),
+		filepath.Join(tempDir, "manifest.json"),
+	})
+	manifestContent, err = os.ReadFile(filepath.Join(tempDir, "manifest.json"))
+	require.NoError(t, err)
+	assert.Equal(t, string(manifestContent), `{
+    "inline-namespace1/my-configmap.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "inline-namespace1/my-deployment.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "inline-namespace1/my-service.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "inline-namespace2/my-deployment.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "inline-namespace2/my-service.yaml": "test-export-envs/inline-envs/main.jsonnet",
+    "static/updated-deployment.yaml": "test-export-envs/static-env/main.jsonnet",
+    "static/updated-service.yaml": "test-export-envs/static-env/main.jsonnet"
+}`)
+}
+
+func BenchmarkExportEnvironmentsWithReplaceEnvs(b *testing.B) {
+	log.SetOutput(io.Discard)
+	tempDir := b.TempDir()
+	require.NoError(b, os.Chdir("testdata"))
+	defer func() { require.NoError(b, os.Chdir("..")) }()
+
+	// Find envs
+	envs, err := FindEnvs("test-export-envs", FindOpts{Selector: labels.Everything()})
+	require.NoError(b, err)
+
+	// Export all envs
+	opts := &ExportEnvOpts{
+		Format:        "{{.metadata.namespace}}/{{.metadata.name}}",
+		Extension:     "yaml",
+		MergeStrategy: ExportMergeStrategyReplaceEnvs,
+	}
+	opts.Opts.ExtCode = jsonnet.InjectedCode{
+		"deploymentName": "'initial-deployment'",
+		"serviceName":    "'initial-service'",
+	}
+	// Export a first time so that the benchmark loops are identical
+	require.NoError(b, ExportEnvironments(envs, tempDir, opts))
+
+	// On every loop, delete manifests from previous envs + reexport all envs
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		require.NoError(b, ExportEnvironments(envs, tempDir, opts), "failed on iteration %d", i)
+	}
+}
+
+func checkFiles(t testing.TB, dir string, files []string) {
+	t.Helper()
+
+	var existingFiles []string
+	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		if info.IsDir() {
+			return nil
+		}
+		existingFiles = append(existingFiles, path)
+		return nil
+	})
+	require.NoError(t, err)
+
+	assert.ElementsMatch(t, files, existingFiles)
+}