refactor: context usage in k8s code (#2405)

## Description refactor context usage in k8s code ## Related Issue Relates to #2365 ## Type of change - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [x] Other (security config, docs update, etc) ## Checklist before merging - [x] Test, docs, adr added or updated as needed - [x] [Contributor Guide Steps](https://github.com/defenseunicorns/zarf/blob/main/CONTRIBUTING.md#developer-workflow) followed --------- Co-authored-by: Austin Abro <37223396+AustinAbro321@users.noreply.github.com>
zarf-dev · May 15, 2024 · 53465d7 · 53465d7
1 parent 48c11d9
commit 53465d7
Show file tree

Hide file tree

Showing 56 changed files with 768 additions and 629 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -9,6 +9,7 @@ repos:
         args:
           - "--allow-missing-credentials"
       - id: detect-private-key
+        exclude: "src/test/e2e/30_config_file_test.go"
       - id: end-of-file-fixer
         exclude: site/src/content/docs/commands/.*
       - id: fix-byte-order-marker

diff --git a/src/cmd/common/utils.go b/src/cmd/common/utils.go
@@ -5,11 +5,13 @@
 package common
 
 import (
+	"context"
 	"os"
 	"os/signal"
 	"syscall"
 
 	"github.com/defenseunicorns/zarf/src/config/lang"
+	"github.com/defenseunicorns/zarf/src/pkg/cluster"
 	"github.com/defenseunicorns/zarf/src/pkg/message"
 )
 
@@ -35,3 +37,14 @@ func ExitOnInterrupt() {
 		}
 	}()
 }
+
+// NewClusterOrDie creates a new Cluster instance and waits for the cluster to be ready or throws a fatal error.
+func NewClusterOrDie(ctx context.Context) *cluster.Cluster {
+	timeoutCtx, cancel := context.WithTimeout(ctx, cluster.DefaultTimeout)
+	defer cancel()
+	c, err := cluster.NewClusterWithWait(timeoutCtx)
+	if err != nil {
+		message.Fatalf(err, "Failed to connect to cluster")
+	}
+	return c
+}
diff --git a/src/cmd/connect.go b/src/cmd/connect.go
@@ -32,7 +32,7 @@ var (
 		Aliases: []string{"c"},
 		Short:   lang.CmdConnectShort,
 		Long:    lang.CmdConnectLong,
-		Run: func(_ *cobra.Command, args []string) {
+		Run: func(cmd *cobra.Command, args []string) {
 			var target string
 			if len(args) > 0 {
 				target = args[0]
@@ -43,12 +43,14 @@ var (
 				spinner.Fatalf(err, lang.CmdConnectErrCluster, err.Error())
 			}
 
+			ctx := cmd.Context()
+
 			var tunnel *k8s.Tunnel
 			if connectResourceName != "" {
 				zt := cluster.NewTunnelInfo(connectNamespace, connectResourceType, connectResourceName, "", connectLocalPort, connectRemotePort)
-				tunnel, err = c.ConnectTunnelInfo(zt)
+				tunnel, err = c.ConnectTunnelInfo(ctx, zt)
 			} else {
-				tunnel, err = c.Connect(target)
+				tunnel, err = c.Connect(ctx, target)
 			}
 			if err != nil {
 				spinner.Fatalf(err, lang.CmdConnectErrService, err.Error())
@@ -90,8 +92,11 @@ var (
 		Use:     "list",
 		Aliases: []string{"l"},
 		Short:   lang.CmdConnectListShort,
-		Run: func(_ *cobra.Command, _ []string) {
-			cluster.NewClusterOrDie().PrintConnectTable()
+		Run: func(cmd *cobra.Command, _ []string) {
+			ctx := cmd.Context()
+			if err := common.NewClusterOrDie(ctx).PrintConnectTable(ctx); err != nil {
+				message.Fatal(err, err.Error())
+			}
 		},
 	}
 )

diff --git a/src/cmd/destroy.go b/src/cmd/destroy.go
@@ -10,10 +10,10 @@ import (
 	"regexp"
 
 	"github.com/defenseunicorns/pkg/helpers"
+	"github.com/defenseunicorns/zarf/src/cmd/common"
 	"github.com/defenseunicorns/zarf/src/config"
 	"github.com/defenseunicorns/zarf/src/config/lang"
 	"github.com/defenseunicorns/zarf/src/internal/packager/helm"
-	"github.com/defenseunicorns/zarf/src/pkg/cluster"
 	"github.com/defenseunicorns/zarf/src/pkg/message"
 	"github.com/defenseunicorns/zarf/src/pkg/utils/exec"
 
@@ -28,16 +28,14 @@ var destroyCmd = &cobra.Command{
 	Aliases: []string{"d"},
 	Short:   lang.CmdDestroyShort,
 	Long:    lang.CmdDestroyLong,
-	Run: func(_ *cobra.Command, _ []string) {
-		c, err := cluster.NewClusterWithWait(cluster.DefaultTimeout)
-		if err != nil {
-			message.Fatalf(err, lang.ErrNoClusterConnection)
-		}
+	Run: func(cmd *cobra.Command, _ []string) {
+		ctx := cmd.Context()
+		c := common.NewClusterOrDie(ctx)
 
 		// NOTE: If 'zarf init' failed to deploy the k3s component (or if we're looking at the wrong kubeconfig)
 		//       there will be no zarf-state to load and the struct will be empty. In these cases, if we can find
 		//       the scripts to remove k3s, we will still try to remove a locally installed k3s cluster
-		state, err := c.LoadZarfState()
+		state, err := c.LoadZarfState(ctx)
 		if err != nil {
 			message.WarnErr(err, lang.ErrLoadState)
 		}
@@ -74,10 +72,12 @@ var destroyCmd = &cobra.Command{
 			helm.Destroy(removeComponents)
 
 			// If Zarf didn't deploy the cluster, only delete the ZarfNamespace
-			c.DeleteZarfNamespace()
+			if err := c.DeleteZarfNamespace(ctx); err != nil {
+				message.Fatal(err, err.Error())
+			}
 
 			// Remove zarf agent labels and secrets from namespaces Zarf doesn't manage
-			c.StripZarfLabelsAndSecretsFromNamespaces()
+			c.StripZarfLabelsAndSecretsFromNamespaces(ctx)
 		}
 	},
 }

diff --git a/src/cmd/dev.go b/src/cmd/dev.go
@@ -40,7 +40,7 @@ var devDeployCmd = &cobra.Command{
 	Args:  cobra.MaximumNArgs(1),
 	Short: lang.CmdDevDeployShort,
 	Long:  lang.CmdDevDeployLong,
-	Run: func(_ *cobra.Command, args []string) {
+	Run: func(cmd *cobra.Command, args []string) {
 		pkgConfig.CreateOpts.BaseDir = common.SetBaseDirectory(args)
 
 		v := common.GetViper()
@@ -50,12 +50,10 @@ var devDeployCmd = &cobra.Command{
 		pkgConfig.PkgOpts.SetVariables = helpers.TransformAndMergeMap(
 			v.GetStringMapString(common.VPkgDeploySet), pkgConfig.PkgOpts.SetVariables, strings.ToUpper)
 
-		// Configure the packager
 		pkgClient := packager.NewOrDie(&pkgConfig)
 		defer pkgClient.ClearTempPaths()
 
-		// Create the package
-		if err := pkgClient.DevDeploy(); err != nil {
+		if err := pkgClient.DevDeploy(cmd.Context()); err != nil {
 			message.Fatalf(err, lang.CmdDevDeployErr, err.Error())
 		}
 	},
@@ -209,19 +207,15 @@ var devFindImagesCmd = &cobra.Command{
 	Run: func(_ *cobra.Command, args []string) {
 		pkgConfig.CreateOpts.BaseDir = common.SetBaseDirectory(args)
 
-		// Ensure uppercase keys from viper
 		v := common.GetViper()
 
 		pkgConfig.CreateOpts.SetVariables = helpers.TransformAndMergeMap(
 			v.GetStringMapString(common.VPkgCreateSet), pkgConfig.CreateOpts.SetVariables, strings.ToUpper)
 		pkgConfig.PkgOpts.SetVariables = helpers.TransformAndMergeMap(
 			v.GetStringMapString(common.VPkgDeploySet), pkgConfig.PkgOpts.SetVariables, strings.ToUpper)
-
-		// Configure the packager
 		pkgClient := packager.NewOrDie(&pkgConfig)
 		defer pkgClient.ClearTempPaths()
 
-		// Find all the images the package might need
 		if _, err := pkgClient.FindImages(); err != nil {
 			message.Fatalf(err, lang.CmdDevFindImagesErr, err.Error())
 		}
@@ -292,8 +286,14 @@ func init() {
 	// use the package create config for this and reset it here to avoid overwriting the config.CreateOptions.SetVariables
 	devFindImagesCmd.Flags().StringToStringVar(&pkgConfig.CreateOpts.SetVariables, "set", v.GetStringMapString(common.VPkgCreateSet), lang.CmdDevFlagSet)
 
-	devFindImagesCmd.Flags().MarkDeprecated("set", "this field is replaced by create-set")
-	devFindImagesCmd.Flags().MarkHidden("set")
+	err := devFindImagesCmd.Flags().MarkDeprecated("set", "this field is replaced by create-set")
+	if err != nil {
+		message.Fatal(err, err.Error())
+	}
+	err = devFindImagesCmd.Flags().MarkHidden("set")
+	if err != nil {
+		message.Fatal(err, err.Error())
+	}
 	devFindImagesCmd.Flags().StringVarP(&pkgConfig.CreateOpts.Flavor, "flavor", "f", v.GetString(common.VPkgCreateFlavor), lang.CmdPackageCreateFlagFlavor)
 	devFindImagesCmd.Flags().StringToStringVar(&pkgConfig.CreateOpts.SetVariables, "create-set", v.GetStringMapString(common.VPkgCreateSet), lang.CmdDevFlagSet)
 	devFindImagesCmd.Flags().StringToStringVar(&pkgConfig.PkgOpts.SetVariables, "deploy-set", v.GetStringMapString(common.VPkgDeploySet), lang.CmdPackageDeployFlagSet)
@@ -341,7 +341,16 @@ func bindDevGenerateFlags(_ *viper.Viper) {
 	generateFlags.StringVar(&pkgConfig.GenerateOpts.Output, "output-directory", "", "Output directory for the generated zarf.yaml")
 	generateFlags.StringVar(&pkgConfig.FindImagesOpts.KubeVersionOverride, "kube-version", "", lang.CmdDevFlagKubeVersion)
 
-	devGenerateCmd.MarkFlagRequired("url")
-	devGenerateCmd.MarkFlagRequired("version")
-	devGenerateCmd.MarkFlagRequired("output-directory")
+	err := devGenerateCmd.MarkFlagRequired("url")
+	if err != nil {
+		message.Fatal(err, err.Error())
+	}
+	err = devGenerateCmd.MarkFlagRequired("version")
+	if err != nil {
+		message.Fatal(err, err.Error())
+	}
+	err = devGenerateCmd.MarkFlagRequired("output-directory")
+	if err != nil {
+		message.Fatal(err, err.Error())
+	}
 }
diff --git a/src/cmd/initialize.go b/src/cmd/initialize.go
@@ -35,7 +35,7 @@ var initCmd = &cobra.Command{
 	Short:   lang.CmdInitShort,
 	Long:    lang.CmdInitLong,
 	Example: lang.CmdInitExample,
-	Run: func(_ *cobra.Command, _ []string) {
+	Run: func(cmd *cobra.Command, _ []string) {
 		zarfLogo := message.GetLogo()
 		_, _ = fmt.Fprintln(os.Stderr, zarfLogo)
 
@@ -58,17 +58,16 @@ var initCmd = &cobra.Command{
 			message.Fatal(err, err.Error())
 		}
 
-		// Ensure uppercase keys from viper
 		v := common.GetViper()
 		pkgConfig.PkgOpts.SetVariables = helpers.TransformAndMergeMap(
 			v.GetStringMapString(common.VPkgDeploySet), pkgConfig.PkgOpts.SetVariables, strings.ToUpper)
 
-		// Configure the packager
 		pkgClient := packager.NewOrDie(&pkgConfig, packager.WithSource(src))
 		defer pkgClient.ClearTempPaths()
 
-		// Deploy everything
-		err = pkgClient.Deploy()
+		ctx := cmd.Context()
+
+		err = pkgClient.Deploy(ctx)
 		if err != nil {
 			message.Fatal(err, err.Error())
 		}

diff --git a/src/cmd/internal.go b/src/cmd/internal.go
@@ -16,7 +16,6 @@ import (
 	"github.com/defenseunicorns/zarf/src/config/lang"
 	"github.com/defenseunicorns/zarf/src/internal/agent"
 	"github.com/defenseunicorns/zarf/src/internal/packager/git"
-	"github.com/defenseunicorns/zarf/src/pkg/cluster"
 	"github.com/defenseunicorns/zarf/src/pkg/message"
 	"github.com/defenseunicorns/zarf/src/types"
 	"github.com/invopop/jsonschema"
@@ -194,15 +193,17 @@ var createReadOnlyGiteaUser = &cobra.Command{
 	Use:   "create-read-only-gitea-user",
 	Short: lang.CmdInternalCreateReadOnlyGiteaUserShort,
 	Long:  lang.CmdInternalCreateReadOnlyGiteaUserLong,
-	Run: func(_ *cobra.Command, _ []string) {
+	Run: func(cmd *cobra.Command, _ []string) {
+		ctx := cmd.Context()
+
 		// Load the state so we can get the credentials for the admin git user
-		state, err := cluster.NewClusterOrDie().LoadZarfState()
+		state, err := common.NewClusterOrDie(ctx).LoadZarfState(ctx)
 		if err != nil {
 			message.WarnErr(err, lang.ErrLoadState)
 		}
 
 		// Create the non-admin user
-		if err = git.New(state.GitServer).CreateReadOnlyUser(); err != nil {
+		if err = git.New(state.GitServer).CreateReadOnlyUser(ctx); err != nil {
 			message.WarnErr(err, lang.CmdInternalCreateReadOnlyGiteaUserErr)
 		}
 	},
@@ -212,24 +213,26 @@ var createPackageRegistryToken = &cobra.Command{
 	Use:   "create-artifact-registry-token",
 	Short: lang.CmdInternalArtifactRegistryGiteaTokenShort,
 	Long:  lang.CmdInternalArtifactRegistryGiteaTokenLong,
-	Run: func(_ *cobra.Command, _ []string) {
-		// Load the state so we can get the credentials for the admin git user
-		c := cluster.NewClusterOrDie()
-		state, err := c.LoadZarfState()
+	Run: func(cmd *cobra.Command, _ []string) {
+		ctx := cmd.Context()
+		c := common.NewClusterOrDie(ctx)
+		state, err := c.LoadZarfState(ctx)
 		if err != nil {
 			message.WarnErr(err, lang.ErrLoadState)
 		}
 
 		// If we are setup to use an internal artifact server, create the artifact registry token
 		if state.ArtifactServer.InternalServer {
-			token, err := git.New(state.GitServer).CreatePackageRegistryToken()
+			token, err := git.New(state.GitServer).CreatePackageRegistryToken(ctx)
 			if err != nil {
 				message.WarnErr(err, lang.CmdInternalArtifactRegistryGiteaTokenErr)
 			}
 
 			state.ArtifactServer.PushToken = token.Sha1
 
-			c.SaveZarfState(state)
+			if err := c.SaveZarfState(ctx, state); err != nil {
+				message.Fatal(err, err.Error())
+			}
 		}
 	},
 }
@@ -238,10 +241,11 @@ var updateGiteaPVC = &cobra.Command{
 	Use:   "update-gitea-pvc",
 	Short: lang.CmdInternalUpdateGiteaPVCShort,
 	Long:  lang.CmdInternalUpdateGiteaPVCLong,
-	Run: func(_ *cobra.Command, _ []string) {
+	Run: func(cmd *cobra.Command, _ []string) {
+		ctx := cmd.Context()
 
 		// There is a possibility that the pvc does not yet exist and Gitea helm chart should create it
-		helmShouldCreate, err := git.UpdateGiteaPVC(rollback)
+		helmShouldCreate, err := git.UpdateGiteaPVC(ctx, rollback)
 		if err != nil {
 			message.WarnErr(err, lang.CmdInternalUpdateGiteaPVCErr)
 		}
@@ -294,6 +298,9 @@ func addHiddenDummyFlag(cmd *cobra.Command, flagDummy string) {
 	if cmd.PersistentFlags().Lookup(flagDummy) == nil {
 		var dummyStr string
 		cmd.PersistentFlags().StringVar(&dummyStr, flagDummy, "", "")
-		cmd.PersistentFlags().MarkHidden(flagDummy)
+		err := cmd.PersistentFlags().MarkHidden(flagDummy)
+		if err != nil {
+			message.Fatal(err, err.Error())
+		}
 	}
 }