fix: Support env var substitution in YAML config (#1041)

Co-authored-by: Kemal Hadimli <disq@users.noreply.github.com>
cloudquery · Jun 29, 2022 · 3bc8014 · 3bc8014
1 parent 867e849
commit 3bc8014
Show file tree

Hide file tree

Showing 11 changed files with 87 additions and 28 deletions.
diff --git a/pkg/config/config_parser.go b/pkg/config/config_parser.go
@@ -1,11 +1,11 @@
 package config
 
 import (
-	"bytes"
 	_ "embed"
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"path/filepath"
 	"strings"
 
@@ -46,7 +46,8 @@ var configSchemaHCL = &hcl.BodySchema{
 
 func (p *Parser) LoadConfigFromSource(name string, data []byte) (*Config, diag.Diagnostics) {
 	if IsNameYAML(name) {
-		return decodeConfigYAML(bytes.NewBuffer(data))
+		newData := os.Expand(string(data), p.getVariableValue)
+		return decodeConfigYAML(strings.NewReader(newData))
 	}
 
 	body, diags := p.LoadFromSource(name, data)

diff --git a/pkg/config/fixtures/config_with_alias.hcl b/pkg/config/fixtures/config_with_alias.hcl
@@ -10,7 +10,7 @@ cloudquery {
 
 provider "aws" {
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "12312312"
     }
     account "ron" {}
@@ -21,10 +21,10 @@ provider "aws" {
 provider "aws" {
   alias = "another-aws"
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "12312312"
     }
-    account "ron" {}
+    accounts "ron" {}
   }
   resources = ["slow_resource"]
 }
diff --git a/pkg/config/fixtures/duplicate_provider_alias.hcl b/pkg/config/fixtures/duplicate_provider_alias.hcl
@@ -11,21 +11,21 @@ cloudquery {
 provider "aws" {
   alias = "same-aws"
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "12312312"
     }
-    account "ron" {}
+    accounts "ron" {}
   }
   resources = ["slow_resource"]
 }
 
 provider "aws" {
   alias = "same-aws"
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "12312312"
     }
-    account "ron" {}
+    accounts "ron" {}
   }
   resources = ["slow_resource"]
 }
diff --git a/pkg/config/fixtures/duplicate_provider_name.hcl b/pkg/config/fixtures/duplicate_provider_name.hcl
@@ -10,20 +10,20 @@ cloudquery {
 
 provider "aws" {
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "12312312"
     }
-    account "ron" {}
+    accounts "ron" {}
   }
   resources = ["slow_resource"]
 }
 
 provider "aws" {
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "12312312"
     }
-    account "ron" {}
+    accounts "ron" {}
   }
   resources = ["slow_resource"]
 }
diff --git a/pkg/config/fixtures/env_vars.hcl b/pkg/config/fixtures/env_vars.hcl
@@ -10,10 +10,10 @@ cloudquery {
 
 provider "aws" {
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "${ROLE_ARN}"
     }
-    account "ron" {}
+    accounts "ron" {}
   }
   resources = ["slow_resource"]
 }
diff --git a/pkg/config/fixtures/env_vars.yml b/pkg/config/fixtures/env_vars.yml
@@ -0,0 +1,16 @@
+cloudquery:
+  connection:
+    dsn: "${DSN}"
+  providers:
+    - name: test
+      source: cloudquery
+      version: v0.0.0
+providers:
+  - name: aws
+    configuration:
+      accounts:
+        - id: "dev"
+          role_arn: "${ROLE_ARN}"
+        - id: "ron"
+    resources:
+      - slow_resource
diff --git a/pkg/config/fixtures/no_source.hcl b/pkg/config/fixtures/no_source.hcl
@@ -9,10 +9,10 @@ cloudquery {
 
 provider "test" {
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "12312312"
     }
-    account "ron" {}
+    accounts "ron" {}
   }
   resources = ["slow_resource"]
 }
diff --git a/pkg/config/fixtures/valid_config.hcl b/pkg/config/fixtures/valid_config.hcl
@@ -10,10 +10,10 @@ cloudquery {
 
 provider "aws" {
   configuration {
-    account "dev" {
+    accounts "dev" {
       role_arn = "12312312"
     }
-    account "ron" {}
+    accounts "ron" {}
   }
   resources = ["slow_resource"]
 }
diff --git a/pkg/config/loader.go b/pkg/config/loader.go
@@ -32,6 +32,8 @@ type Parser struct {
 	fs         afero.Afero
 	p          *hclparse.Parser
 	HCLContext hcl.EvalContext
+
+	variables map[string]string
 }
 
 type Option func(*Parser)
@@ -48,6 +50,17 @@ func WithFS(aferoFs afero.Fs) Option {
 func WithEnvironmentVariables(prefix string, vars []string) Option {
 	return func(p *Parser) {
 		EnvToHCLContext(&p.HCLContext, prefix, vars)
+
+		for _, v := range vars {
+			pair := strings.SplitN(v, "=", 2)
+			if strings.HasPrefix(pair[0], prefix) {
+				var varVal string
+				if len(pair) == 2 {
+					varVal = pair[1]
+				}
+				p.variables[strings.TrimPrefix(pair[0], prefix)] = varVal
+			}
+		}
 	}
 }
 
@@ -67,6 +80,7 @@ func NewParser(options ...Option) *Parser {
 			Variables: make(map[string]cty.Value),
 			Functions: make(map[string]function.Function),
 		},
+		variables: make(map[string]string),
 	}
 
 	for _, opt := range options {
@@ -75,6 +89,10 @@ func NewParser(options ...Option) *Parser {
 	return &p
 }
 
+func (p *Parser) getVariableValue(s string) string {
+	return p.variables[s]
+}
+
 // LoadFile is a low-level method that reads the file at the given path
 func (p *Parser) LoadFile(path string) ([]byte, diag.Diagnostics) {
 	var contents []byte

diff --git a/pkg/config/loader_test.go b/pkg/config/loader_test.go
@@ -17,6 +17,7 @@ import (
 	"github.com/johannesboyne/gofakes3/backend/s3mem"
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/assert"
+	"gopkg.in/yaml.v3"
 )
 
 const testConfig = `cloudquery {
@@ -48,16 +49,16 @@ const bucketName = "myBucket"
 const defaultPermissions = 0644
 
 type Account struct {
-	ID      string `hcl:",label"`
-	RoleARN string `hcl:"role_arn,optional"`
+	ID      string `yaml:"id" hcl:",label"`
+	RoleARN string `yaml:"role_arn,omitempty" hcl:"role_arn,optional"`
 }
 
 type AwsConfig struct {
-	Regions    []string  `hcl:"regions,optional"`
-	Accounts   []Account `hcl:"account,block"`
-	AWSDebug   bool      `hcl:"aws_debug,optional"`
-	MaxRetries int       `hcl:"max_retries,optional" default:"5"`
-	MaxBackoff int       `hcl:"max_backoff,optional" default:"30"`
+	Regions    []string  `yaml:"regions,omitempty" hcl:"regions,optional"`
+	Accounts   []Account `yaml:"accounts" hcl:"accounts,block"`
+	AWSDebug   bool      `yaml:"aws_debug,omitempty" hcl:"aws_debug,optional"`
+	MaxRetries int       `yaml:"max_retries,omitempty" hcl:"max_retries,optional" default:"5"`
+	MaxBackoff int       `yaml:"max_backoff,omitempty" hcl:"max_backoff,optional" default:"30"`
 }
 
 func TestParser_LoadValidConfigFromFile(t *testing.T) {
@@ -145,7 +146,7 @@ func TestWithEnvironmentVariables(t *testing.T) {
 	assert.Equal(t, "value 2", p.HCLContext.Variables["Var2"].AsString())
 }
 
-func TestConfigEnvVariableSubstitution(t *testing.T) {
+func TestConfigEnvVariableSubstitutionHCL(t *testing.T) {
 	p := NewParser(WithEnvironmentVariables(EnvVarPrefix, []string{
 		"CQ_VAR_DSN=postgres://postgres:pass@localhost:5432/postgres",
 		"CQ_VAR_ROLE_ARN=12312312",
@@ -166,6 +167,27 @@ func TestConfigEnvVariableSubstitution(t *testing.T) {
 	assert.Equal(t, "12312312", c.Accounts[0].RoleARN)
 }
 
+func TestConfigEnvVariableSubstitutionYAML(t *testing.T) {
+	p := NewParser(WithEnvironmentVariables(EnvVarPrefix, []string{
+		"CQ_VAR_DSN=postgres://postgres:pass@localhost:5432/postgres",
+		"CQ_VAR_ROLE_ARN=12312312",
+	}))
+	cfg, diags := p.LoadConfigFile("fixtures/env_vars.yml")
+	if diags != nil {
+		for _, d := range diags {
+			t.Error(d.Error())
+		}
+		return
+	}
+	assert.Equal(t, "postgres://postgres:pass@localhost:5432/postgres", cfg.CloudQuery.Connection.DSN)
+
+	c := AwsConfig{}
+	errs := yaml.Unmarshal(cfg.Providers[0].Configuration, &c)
+	assert.Nil(t, errs)
+
+	assert.Equal(t, "12312312", c.Accounts[0].RoleARN)
+}
+
 func TestLoader_LoadConfigNoSourceField(t *testing.T) {
 	p := NewParser()
 	cfg, diags := p.LoadConfigFile("fixtures/no_source.hcl")

diff --git a/pkg/core/core_test.go b/pkg/core/core_test.go
@@ -43,7 +43,9 @@ func setupDB(t *testing.T) (dsn string) {
 		}
 
 		if _, err := conn.Exec(context.Background(), "DROP DATABASE "+newDB+" WITH(FORCE)"); err != nil {
-			t.Logf("teardown: drop database failed: %v", err)
+			if _, err := conn.Exec(context.Background(), "DROP DATABASE "+newDB); err != nil {
+				t.Logf("teardown: drop database failed: %v", err)
+			}
 		}
 	})