Added unset-not-found flag

README.md

@@ -29,7 +29,7 @@ $ aws ssm put-parameter --name /testing/my-app/dbpass --value "some-secret-passw
 $ aws ssm put-parameter --name /testing/my-app/privatekey --value "some-private-key" --type SecureString --key-id "alias/aws/ssm" --region us-east-1
 ```
 
-2. Install aws-env using static binary (amd64 only) (choose proper [version](https://github.com/sendgrid/aws-env/releases)). 
+2. Install aws-env using static binary (amd64 only) (choose proper [version](https://github.com/sendgrid/aws-env/releases)).
 
 ```
 $ wget https://github.com/sendgrid/aws-env/releases/download/1.4.0/aws-env -O aws-env
@@ -118,20 +118,20 @@ initializing and passing an aws-sdk-go-v2 SSM client.
 
 ### Use to update a file in-place
 
-The `-f` flag can be used to pass in a file to update in-place rather than 
-operating on the environment variables. It will only update the first 
+The `-f` flag can be used to pass in a file to update in-place rather than
+operating on the environment variables. It will only update the first
 occurrence per line. It stops parsing when a character is no longer a valid
-Parameter Store path. 
+Parameter Store path.
 
 Example usage
 ```
-$ mrroboto upload -p /path/to/the/username -v localtestuser               
+$ mrroboto upload -p /path/to/the/username -v localtestuser
 2019/03/13 14:15:04 parameter=/path/to/the/username regions=[us-east-1 us-east-2 us-west-1 us-west-2]
 
-$ mrroboto upload -p /path/to/the/password -v localtestpass            
+$ mrroboto upload -p /path/to/the/password -v localtestpass
 2019/03/13 14:15:14 parameter=/path/to/the/password regions=[us-east-1 us-east-2 us-west-1 us-west-2]
 
-$ cat test.txt4                                                           
+$ cat test.txt4
 mysql_users:
  (
     {
@@ -144,11 +144,11 @@ mysql_users:
     }
  )
 
-$ ./aws-env -f test.txt4                                                  
+$ ./aws-env -f test.txt4
 INFO[0000] aws-env starting                              app_version=0.0.1 built_at="Wed Mar 13 20:12:42 UTC 2019" git_hash=545515b6b6646f1bd2f95dea13478066782deb0c
 
 
-$ cat test.txt4                                                          
+$ cat test.txt4
 mysql_users:
  (
     {
@@ -204,6 +204,12 @@ aws-env exposes an `--assume-role` flag (or `AWS_ENV_ASSUME_ROLE`). This can
 be used to further assume roles if you have to gain access using a chain of
 roles.
 
+## Unset not found variables
+The default behaviour is to raise error in case some of the mapped variables
+are not found in AWS. You can pass --unset-not-found flaf (or `AWS_ENV_UNSET_NOT_FOUND`)
+to force the variable to present in AWS to be unset. Note that this does not
+work with combination with file replacement.
+
 ### Example Assume Role
 In Kubernetes, if you are using Annotations with a service role, `kube2iam`
 will assume your service role using the metadata service. You can then use

awsenv/helpers.go

@@ -84,3 +84,22 @@ func pathmap(prefix string, env []string) map[string]string {
 
 	return m
 }
+
+func unmappedVars(prefix string, env []string) []string {
+	var vars []string
+
+	for _, rawVar := range env {
+		idx := strings.Index(rawVar, "=")
+		if idx < 0 {
+			// impossible on real systems?
+			continue
+		}
+
+		name, path := rawVar[:idx], rawVar[idx+1:]
+		if strings.HasPrefix(path, prefix) {
+			vars = append(vars, name)
+		}
+	}
+
+	return vars
+}

awsenv/replacer.go

@@ -2,15 +2,16 @@ package awsenv
 
 import (
 	"context"
+	"fmt"
 	"os"
 
-	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
 )
 
 var (
-	setenv  = os.Setenv
-	environ = os.Environ
+	setenv   = os.Setenv
+	unsetenv = os.Unsetenv
+	environ  = os.Environ
 )
 
 // DefaultPrefix holds the standard environment value prefix.
@@ -35,22 +36,24 @@ type LimitedParamsGetter interface {
 // given value prefix, using the given ParamsGetter.
 //
 // NewReplacer will panic if envValuePrefix is the empty string.
-func NewReplacer(envValuePrefix string, ssm ParamsGetter) *Replacer {
+func NewReplacer(envValuePrefix string, unsetNotFound bool, ssm ParamsGetter) *Replacer {
 	if envValuePrefix == "" {
 		panic("awsenv: envValuePrefix must be non-empty")
 	}
 
 	return &Replacer{
-		ssm:    ssm,
-		prefix: envValuePrefix,
+		unsetNotFound: unsetNotFound,
+		ssm:           ssm,
+		prefix:        envValuePrefix,
 	}
 }
 
 // Replacer handles replacing existing environment variables with values
 // retrieved from AWS Parameter Store.
 type Replacer struct {
-	ssm    ParamsGetter
-	prefix string
+	unsetNotFound bool
+	ssm           ParamsGetter
+	prefix        string
 }
 
 // ReplaceAll overwrites applicable environment variables with values
@@ -70,6 +73,17 @@ func (r *Replacer) ReplaceAll(ctx context.Context) error {
 		}
 	}
 
+	unmapped := unmappedVars(r.prefix, environ())
+	if len(unmapped) > 0 && !r.unsetNotFound {
+		return fmt.Errorf("following variables not found in AWS: %v", unmapped)
+	}
+	for _, name := range unmapped {
+		suberr := unsetenv(name)
+		if err == nil && suberr != nil {
+			err = suberr
+		}
+	}
+
 	return err
 }
 
@@ -110,7 +124,6 @@ func fetch(ctx context.Context, ssm ParamsGetter, paths []string) (map[string]st
 		// copied to avoid race condition
 		i := i
 		batch := batches[i]
-
 		eg.Go(func() error {
 			var err error
 			results[i], err = ssm.GetParams(ctx, batch)
@@ -127,12 +140,5 @@ func fetch(ctx context.Context, ssm ParamsGetter, paths []string) (map[string]st
 	dest := make(map[string]string, len(paths))
 	merge(dest, results)
 
-	for _, path := range paths {
-		_, ok := dest[path]
-		if !ok {
-			return dest, errors.Errorf("awsenv: param not found: %q", path)
-		}
-	}
-
 	return dest, nil
 }

awsenv/replacer_test.go

@@ -13,7 +13,7 @@ func TestReplacer_panic(t *testing.T) {
 		return nil, errors.New("no implementation")
 	})
 
-	require.Panics(t, func() { NewReplacer("", mockGetter) })
+	require.Panics(t, func() { NewReplacer("", false, mockGetter) })
 }
 
 func TestReplacer_ReplaceAll_noop(t *testing.T) {
@@ -25,7 +25,7 @@ func TestReplacer_ReplaceAll_noop(t *testing.T) {
 	env.install()
 
 	ctx := context.Background()
-	r := NewReplacer("awsenv:", mockGetter)
+	r := NewReplacer("awsenv:", false, mockGetter)
 	err := r.ReplaceAll(ctx)
 	require.NoError(t, err, "expected no error")
 	require.Empty(t, env)
@@ -44,7 +44,7 @@ func TestReplacerMultiple(t *testing.T) {
 		"/param/path/here/v2": "val2",
 	}
 
-	r := NewReplacer(DefaultPrefix, params)
+	r := NewReplacer(DefaultPrefix, false, params)
 
 	ctx := context.Background()
 	err := r.ReplaceAll(ctx)
@@ -69,7 +69,7 @@ func TestReplacerNotFound(t *testing.T) {
 
 	var params mockParamStore
 
-	r := NewReplacer("awsenv:", params)
+	r := NewReplacer("awsenv:", false, params)
 
 	ctx := context.Background()
 	err := r.ReplaceAll(ctx)
@@ -87,13 +87,30 @@ func TestReplacerMissing(t *testing.T) {
 		return nil, nil
 	}
 
-	r := NewReplacer("awsenv:", mockParamsGetter(getter))
+	r := NewReplacer("awsenv:", false, mockParamsGetter(getter))
 	ctx := context.Background()
 
 	err := r.ReplaceAll(ctx)
 	require.Error(t, err, "expected an error")
 }
 
+func TestReplacerMissingWithUnsetNotFound(t *testing.T) {
+	env := fakeEnv{
+		"SOME_SECRET": "awsenv:/param/path/here/doesnt/exist", // match
+	}
+	env.install()
+
+	getter := func(context.Context, []string) (map[string]string, error) {
+		return nil, nil
+	}
+
+	r := NewReplacer("awsenv:", true, mockParamsGetter(getter))
+	ctx := context.Background()
+
+	err := r.ReplaceAll(ctx)
+	require.NoError(t, err, "expected no error")
+}
+
 type mockParamStore map[string]string
 
 func (m mockParamStore) GetParams(ctx context.Context, paths []string) (map[string]string, error) {

cmd/aws-env/main.go

@@ -30,12 +30,13 @@ var (
 
 	app = initApp()
 
-	prefix     string
-	region     string
-	profile    string
-	assumeRole string
-	fileName   string
-	ecs        bool
+	prefix        string
+	unsetNotFound bool
+	region        string
+	profile       string
+	assumeRole    string
+	fileName      string
+	ecs           bool
 )
 
 const description = `
@@ -93,6 +94,12 @@ func initApp() *cli.App {
 			Usage:       "enable ECS mode, using the default credential provider to support ECS",
 			Destination: &ecs,
 		},
+		cli.BoolFlag{
+			Name:        "unset-not-found",
+			EnvVar:      "AWS_ENV_UNSET_NOT_FOUND",
+			Usage:       "Unset variable not found in AWS",
+			Destination: &unsetNotFound,
+		},
 	}
 	newApp.Commands = append(newApp.Commands, cli.Command{
 		Name:   "licenses",
@@ -166,7 +173,7 @@ func run(c *cli.Context) error {
 }
 
 func envReplacement(c *cli.Context, ssmClient *ssm.SSM) error {
-	r := awsenv.NewReplacer(prefix, v1.NewParamsGetter(ssmClient))
+	r := awsenv.NewReplacer(prefix, unsetNotFound, v1.NewParamsGetter(ssmClient))
 
 	if c.NArg() == 0 {
 		return dump(r)