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

…#760)

* Add `--delete-previous` option to `tk export`

Documented in `exporting.md`. This option allows us to delete the previously exported manifests for an environment.

This is useful when exporting a single environment and merging the result with an existing GitOps repository, rather than re-exporting all environments.

I also added benchmark and a test for the export code.

* Fix lint

* Replace with `--merge-strategy`

* PR comments

* Validate possible merge values

* Describe default behavior

* Add test and extract the code. Linting should pass now

cmd/tk/export.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"regexp"
 	"runtime"
@@ -34,12 +35,17 @@ 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. The default setting is to disallow exporting to an existing directory. Values: 'fail-on-conficts', 'replace-envs'")
+
 	vars := workflowFlags(cmd.Flags())
 	getJsonnetOpts := jsonnetFlags(cmd.Flags())
 	getLabelSelector := labelSelectorFlag(cmd.Flags())
@@ -59,7 +65,6 @@ func exportCmd() *cli.Command {
 		opts := tanka.ExportEnvOpts{
 			Format:    *format,
 			Extension: *extension,
-			Merge:     *merge,
 			Opts: tanka.Opts{
 				JsonnetOpts: getJsonnetOpts(),
 				Filters:     filters,
@@ -68,6 +73,11 @@ func exportCmd() *cli.Command {
 			Selector:    getLabelSelector(),
 			Parallelism: *parallel,
 		}
+
+		if opts.MergeStrategy, err = determineMergeStrategy(*merge, *mergeStrategy); err != nil {
+			return err
+		}
+
 		opts.Opts.CachePath = *cachePath
 		for _, expr := range *cacheEnvs {
 			regex, err := regexp.Compile(expr)
@@ -116,3 +126,20 @@ func exportCmd() *cli.Command {
 	}
 	return cmd
 }
+
+// `--merge` is deprecated in favor of `--merge-strategy`. However, merge has to keep working for now.
+func determineMergeStrategy(deprecatedMergeFlag bool, mergeStrategy string) (tanka.ExportMergeStrategy, error) {
+	if deprecatedMergeFlag && mergeStrategy != "" {
+		return "", errors.New("cannot use --merge and --merge-strategy at the same time")
+	}
+	if deprecatedMergeFlag {
+		return tanka.ExportMergeStrategyFailConflicts, nil
+	}
+
+	switch strategy := tanka.ExportMergeStrategy(mergeStrategy); strategy {
+	case tanka.ExportMergeStrategyFailConflicts, tanka.ExportMergeStrategyReplaceEnvs, tanka.ExportMergeStrategyNone:
+		return strategy, nil
+	}
+
+	return "", fmt.Errorf("invalid merge strategy: %q", mergeStrategy)
+}

cmd/tk/export_test.go

@@ -0,0 +1,63 @@
+package main
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/grafana/tanka/pkg/tanka"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDetermineMergeStrategy(t *testing.T) {
+	cases := []struct {
+		name           string
+		deprecatedFlag bool
+		mergeStrategy  string
+		expected       tanka.ExportMergeStrategy
+		expectErr      error
+	}{
+		{
+			name:           "default",
+			deprecatedFlag: false,
+			mergeStrategy:  "",
+			expected:       tanka.ExportMergeStrategyNone,
+		},
+		{
+			name:           "deprecated flag set",
+			deprecatedFlag: true,
+			expected:       tanka.ExportMergeStrategyFailConflicts,
+		},
+		{
+			name:           "both values set",
+			deprecatedFlag: true,
+			mergeStrategy:  "fail-conflicts",
+			expectErr:      errors.New("cannot use --merge and --merge-strategy at the same time"),
+		},
+		{
+			name:          "fail-conflicts",
+			mergeStrategy: "fail-on-conflicts",
+			expected:      tanka.ExportMergeStrategyFailConflicts,
+		},
+		{
+			name:          "replace-envs",
+			mergeStrategy: "replace-envs",
+			expected:      tanka.ExportMergeStrategyReplaceEnvs,
+		},
+		{
+			name:          "bad value",
+			mergeStrategy: "blabla",
+			expectErr:     errors.New("invalid merge strategy: \"blabla\""),
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			result, err := determineMergeStrategy(tc.deprecatedFlag, tc.mergeStrategy)
+			if tc.expectErr != nil {
+				assert.EqualError(t, err, tc.expectErr.Error())
+			} else {
+				assert.NoError(t, err)
+			}
+			assert.Equal(t, tc.expected, result)
+		})
+	}
+}

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)