Repository: ariga/atlas
Pull request #1200 (2022-10-15 14:24:00)

cmd/atlas: add 'sql' data-source to config files


Patch:
diff --git a/cmd/atlas/internal/cmdapi/cmdapi.go b/cmd/atlas/internal/cmdapi/cmdapi.go
index e768a8e8..de6164bb 100644
--- a/cmd/atlas/internal/cmdapi/cmdapi.go
+++ b/cmd/atlas/internal/cmdapi/cmdapi.go
@@ -73,6 +73,13 @@ Atlas is licensed under Apache 2.0 as found in https://github.com/ariga/atlas/bl
 func init() {
 	Root.AddCommand(versionCmd)
 	Root.AddCommand(licenseCmd)
+	// Register a global function to clean up the global
+	// flags regardless if the command passed or failed.
+	cobra.OnFinalize(func() {
+		GlobalFlags.ConfigURL = ""
+		GlobalFlags.Vars = nil
+		GlobalFlags.SelectedEnv = ""
+	})
 }
 
 // inputValsFromEnv populates GlobalFlags.Vars from the active environment. If we are working
diff --git a/cmd/atlas/internal/cmdapi/project.go b/cmd/atlas/internal/cmdapi/project.go
index af32fd2d..1664baba 100644
--- a/cmd/atlas/internal/cmdapi/project.go
+++ b/cmd/atlas/internal/cmdapi/project.go
@@ -5,6 +5,7 @@
 package cmdapi
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/url"
@@ -12,7 +13,10 @@ import (
 	"path/filepath"
 
 	"ariga.io/atlas/schemahcl"
+	"ariga.io/atlas/sql/sqlclient"
 
+	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/hcl/v2/hclsyntax"
 	"github.com/zclconf/go-cty/cty"
 )
 
@@ -176,6 +180,9 @@ func (e *Env) asMap() (map[string]string, error) {
 
 var hclState = schemahcl.New(
 	schemahcl.WithScopedEnums("env.migration.format", formatAtlas, formatFlyway, formatLiquibase, formatGoose, formatGolangMigrate),
+	schemahcl.WithDataSource("sql", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
+		return (&sqlsrc{ctx: ctx, block: b}).exec()
+	}),
 )
 
 // LoadEnv reads the project file in path, and loads
@@ -228,3 +235,106 @@ func LoadEnv(name string, opts ...LoadOption) ([]*Env, error) {
 func init() {
 	schemahcl.Register("env", &Env{})
 }
+
+// sqlsrc represents an SQL data-source.
+type sqlsrc struct {
+	ctx   *hcl.EvalContext
+	block *hclsyntax.Block
+}
+
+// exec executes the source block for getting the data.
+func (s *sqlsrc) exec() (cty.Value, error) {
+	attrs, diags := s.block.Body.JustAttributes()
+	if diags.HasErrors() {
+		return cty.NilVal, diags
+	}
+	u, err := s.stringAttr(attrs, "url")
+	if err != nil {
+		return cty.NilVal, err
+	}
+	query, err := s.stringAttr(attrs, "query")
+	if err != nil {
+		return cty.NilVal, err
+	}
+	var args []any
+	if at, ok := attrs["args"]; ok {
+		switch v, diags := at.Expr.Value(s.ctx); {
+		case diags.HasErrors():
+			return cty.NilVal, s.errorf(`evaluating "args": %w`, err)
+		case !v.CanIterateElements():
+			return cty.NilVal, s.errorf(`attribute "args" must be a list, got: %s`, v.Type())
+		default:
+			for it := v.ElementIterator(); it.Next(); {
+				switch _, v := it.Element(); v.Type() {
+				case cty.String:
+					args = append(args, v.AsString())
+				case cty.Number:
+					f, _ := v.AsBigFloat().Float64()
+					args = append(args, f)
+				case cty.Bool:
+					args = append(args, v.True())
+				default:
+					return cty.NilVal, fmt.Errorf(`attribute "args" must contain primitive types, got: %s`, v.Type())
+				}
+			}
+		}
+	}
+	var values []cty.Value
+	c, err := sqlclient.Open(context.Background(), u)
+	if err != nil {
+		return cty.NilVal, s.errorf("opening connection: %w", err)
+	}
+	defer c.Close()
+	rows, err := c.QueryContext(context.Background(), query, args...)
+	if err != nil {
+		return cty.NilVal, s.errorf("executing query: %w", err)
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var v any
+		if err := rows.Scan(&v); err != nil {
+			return cty.NilVal, s.errorf("scanning row: %w", err)
+		}
+		switch v := v.(type) {
+		case bool:
+			values = append(values, cty.BoolVal(v))
+		case int64:
+			values = append(values, cty.NumberIntVal(v))
+		case float64:
+			values = append(values, cty.NumberFloatVal(v))
+		case string:
+			values = append(values, cty.StringVal(v))
+		default:
+			return cty.NilVal, s.errorf("unsupported row type: %T", v)
+		}
+	}
+	obj := map[string]cty.Value{
+		"count":  cty.NumberIntVal(int64(len(values))),
+		"values": cty.ListValEmpty(cty.NilType),
+		"value":  cty.NilVal,
+	}
+	if len(values) > 0 {
+		obj["value"] = values[0]
+		obj["values"] = cty.ListVal(values)
+	}
+	return cty.ObjectVal(obj), nil
+}
+
+func (s *sqlsrc) stringAttr(attrs hcl.Attributes, name string) (string, error) {
+	at, ok := attrs[name]
+	if !ok {
+		return "", s.errorf("missing %q attribute", name)
+	}
+	u, diags := at.Expr.Value(s.ctx)
+	if diags.HasErrors() {
+		return "", s.errorf("evaluating %q attribute: %w", name, diags)
+	}
+	if u.Type() != cty.String {
+		return "", s.errorf("attribute %q must be a string, got: %s", name, u.Type())
+	}
+	return u.AsString(), nil
+}
+
+func (s *sqlsrc) errorf(format string, args ...any) error {
+	return fmt.Errorf("data.sql.%s: %w", s.block.Labels[1], fmt.Errorf(format, args...))
+}
diff --git a/schemahcl/context.go b/schemahcl/context.go
index bc8f9dbe..5ef2da3b 100644
--- a/schemahcl/context.go
+++ b/schemahcl/context.go
@@ -32,14 +32,11 @@ type blockVar struct {
 //	  default = "rotemtam"
 //	}
 func (s *State) setInputVals(ctx *hcl.EvalContext, body hcl.Body, input map[string]cty.Value) error {
-	var (
-		doc struct {
-			Vars   []*blockVar `hcl:"variable,block"`
-			Remain hcl.Body    `hcl:",remain"`
-		}
-		nctx = varBlockContext(ctx)
-	)
-	if diag := gohcl.DecodeBody(body, nctx, &doc); diag.HasErrors() {
+	var doc struct {
+		Vars   []*blockVar `hcl:"variable,block"`
+		Remain hcl.Body    `hcl:",remain"`
+	}
+	if diag := gohcl.DecodeBody(body, ctx, &doc); diag.HasErrors() {
 		return diag
 	}
 	ctxVars := make(map[string]cty.Value)
@@ -74,7 +71,7 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 	type node struct {
 		addr  [3]string
 		edges func() []hcl.Traversal
-		value func() (cty.Value, hcl.Diagnostics)
+		value func() (cty.Value, error)
 	}
 	var (
 		nodes  = make(map[[3]string]*node)
@@ -95,7 +92,7 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			addr := [3]string{dataBlock, b.Labels[0], b.Labels[1]}
 			nodes[addr] = &node{
 				addr:  addr,
-				value: func() (cty.Value, hcl.Diagnostics) { return h(ctx, b) },
+				value: func() (cty.Value, error) { return h(ctx, b) },
 				edges: func() []hcl.Traversal { return bodyVars(b.Body) },
 			}
 		case localsBlock:
@@ -106,8 +103,14 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 				addr := [3]string{localRef, k, ""}
 				nodes[addr] = &node{
 					addr:  addr,
-					value: func() (cty.Value, hcl.Diagnostics) { return v.Expr.Value(ctx) },
 					edges: func() []hcl.Traversal { return hclsyntax.Variables(v.Expr) },
+					value: func() (cty.Value, error) {
+						v, diags := v.Expr.Value(ctx)
+						if diags.HasErrors() {
+							return cty.NilVal, diags
+						}
+						return v, nil
+					},
 				}
 			}
 		default:
@@ -148,9 +151,9 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			}
 		}
 		delete(progress, n)
-		v, diags := n.value()
-		if diags.HasErrors() {
-			return diags
+		v, err := n.value()
+		if err != nil {
+			return err
 		}
 		switch n.addr[0] {
 		case dataBlock:
diff --git a/schemahcl/opts.go b/schemahcl/opts.go
index 8663bf51..42de4f0c 100644
--- a/schemahcl/opts.go
+++ b/schemahcl/opts.go
@@ -21,7 +21,7 @@ type (
 		types    []*TypeSpec
 		newCtx   func() *hcl.EvalContext
 		pathVars map[string]map[string]cty.Value
-		datasrc  map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, hcl.Diagnostics)
+		datasrc  map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error)
 	}
 
 	// Option configures a Config.
@@ -33,10 +33,10 @@ func New(opts ...Option) *State {
 	cfg := &Config{
 		pathVars: make(map[string]map[string]cty.Value),
 		newCtx: func() *hcl.EvalContext {
-			return &hcl.EvalContext{
+			return stdTypes(&hcl.EvalContext{
 				Functions: stdFuncs(),
 				Variables: make(map[string]cty.Value),
-			}
+			})
 		},
 	}
 	for _, opt := range opts {
@@ -88,10 +88,10 @@ func WithScopedEnums(path string, enums ...string) Option {
 //	data "text" "hello" {
 //	  value = "hello world"
 //	}
-func WithDataSource(name string, h func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, hcl.Diagnostics)) Option {
+func WithDataSource(name string, h func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error)) Option {
 	return func(c *Config) {
 		if c.datasrc == nil {
-			c.datasrc = make(map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, hcl.Diagnostics))
+			c.datasrc = make(map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error))
 		}
 		c.datasrc[name] = h
 	}
@@ -100,10 +100,10 @@ func WithDataSource(name string, h func(*hcl.EvalContext, *hclsyntax.Block) (cty
 // WithTypes configures the list of given types as identifiers in the unmarshaling context.
 func WithTypes(typeSpecs []*TypeSpec) Option {
 	newCtx := func() *hcl.EvalContext {
-		ctx := &hcl.EvalContext{
+		ctx := stdTypes(&hcl.EvalContext{
+			Functions: stdFuncs(),
 			Variables: make(map[string]cty.Value),
-			Functions: make(map[string]function.Function),
-		}
+		})
 		for _, ts := range typeSpecs {
 			typeSpec := ts
 			// If no required args exist, register the type as a variable in the HCL context.
@@ -119,9 +119,9 @@ func WithTypes(typeSpecs []*TypeSpec) Option {
 		ctx.Functions["sql"] = rawExprImpl()
 		return ctx
 	}
-	return func(config *Config) {
-		config.newCtx = newCtx
-		config.types = append(config.types, typeSpecs...)
+	return func(c *Config) {
+		c.newCtx = newCtx
+		c.types = append(c.types, typeSpecs...)
 	}
 }
 
diff --git a/schemahcl/stdlib.go b/schemahcl/stdlib.go
index 83b01733..83ea5646 100644
--- a/schemahcl/stdlib.go
+++ b/schemahcl/stdlib.go
@@ -15,7 +15,7 @@ import (
 	"github.com/zclconf/go-cty/cty/function/stdlib"
 )
 
-func varBlockContext(ctx *hcl.EvalContext) *hcl.EvalContext {
+func stdTypes(ctx *hcl.EvalContext) *hcl.EvalContext {
 	ctx = ctx.NewChild()
 	ctx.Variables = map[string]cty.Value{
 		"string": cty.CapsuleVal(ctyNilType, &cty.String),
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index 964b86f0..018b2f90 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -364,7 +364,7 @@ Top:
 			switch c := s[i]; {
 			case c == '\\', c == quote && s[i+1] == quote:
 				i++
-			// Accept only escaped quotes and eject otherwise.
+			// Accept only escaped quotes and reject otherwise.
 			case c == quote:
 				continue Top
 			}


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate_test.go b/cmd/atlas/internal/cmdapi/migrate_test.go
index ea54a71e..49daf982 100644
--- a/cmd/atlas/internal/cmdapi/migrate_test.go
+++ b/cmd/atlas/internal/cmdapi/migrate_test.go
@@ -393,8 +393,9 @@ func TestMigrate_Apply(t *testing.T) {
 }
 
 func TestMigrate_ApplyMultiEnv(t *testing.T) {
-	p := t.TempDir()
-	h := `
+	t.Run("FromVars", func(t *testing.T) {
+		p := t.TempDir()
+		h := `
 variable "urls" {
   type = list(string)
 }
@@ -408,25 +409,97 @@ env "local" {
   }
 }
 `
-	path := filepath.Join(p, "atlas.hcl")
-	err := os.WriteFile(path, []byte(h), 0600)
-	require.NoError(t, err)
-	cmd := migrateCmd()
-	t.Cleanup(func() { GlobalFlags.SelectedEnv = "" })
-	cmd.AddCommand(migrateApplyCmd())
-	s, err := runCmd(
-		cmd, "apply",
-		"-c", "file://"+path,
-		"--env", "local",
-		"--var", fmt.Sprintf("urls=sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test1.db")),
-		"--var", fmt.Sprintf("urls=sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test2.db")),
-	)
-	require.NoError(t, err)
-	require.Equal(t, 2, strings.Count(s, "Migrating to version 20220318104615 (2 migrations in total)"), "execution per environment")
-	_, err = os.Stat(filepath.Join(p, "test1.db"))
-	require.NoError(t, err)
-	_, err = os.Stat(filepath.Join(p, "test2.db"))
-	require.NoError(t, err)
+		path := filepath.Join(p, "atlas.hcl")
+		err := os.WriteFile(path, []byte(h), 0600)
+		require.NoError(t, err)
+		cmd := migrateCmd()
+		cmd.AddCommand(migrateApplyCmd())
+		s, err := runCmd(
+			cmd, "apply",
+			"-c", "file://"+path,
+			"--env", "local",
+			"--var", fmt.Sprintf("urls=sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test1.db")),
+			"--var", fmt.Sprintf("urls=sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "test2.db")),
+		)
+		require.NoError(t, err)
+		require.Equal(t, 2, strings.Count(s, "Migrating to version 20220318104615 (2 migrations in total)"), "execution per environment")
+		_, err = os.Stat(filepath.Join(p, "test1.db"))
+		require.NoError(t, err)
+		_, err = os.Stat(filepath.Join(p, "test2.db"))
+		require.NoError(t, err)
+	})
+
+	t.Run("FromDataSrc", func(t *testing.T) {
+		var (
+			h = `
+variable "url" {
+  type = string
+}
+
+data "sql" "tenants" {
+  url   = var.url
+  query = "SELECT name FROM tenants"
+}
+
+env "local" {
+  for_each = toset(data.sql.tenants.values)
+  url = "sqlite://file:${each.value}?cache=shared&_fk=1"
+  dev = "sqlite://ci?mode=memory&cache=shared&_fk=1"
+  migration {
+    dir = "file://testdata/sqlite"
+  }
+}
+`
+			p    = t.TempDir()
+			path = filepath.Join(p, "atlas.hcl")
+			dbs  = []string{filepath.Join(p, "test1.db"), filepath.Join(p, "test2.db"), filepath.Join(p, "test3.db")}
+		)
+		err := os.WriteFile(path, []byte(h), 0600)
+		require.NoError(t, err)
+		db, err := sql.Open("sqlite3", fmt.Sprintf("file:%s?cache=shared&_fk=1", filepath.Join(p, "tenants.db")))
+		require.NoError(t, err)
+		_, err = db.Exec("CREATE TABLE `tenants` (`name` TEXT)")
+		require.NoError(t, err)
+		_, err = db.Exec("INSERT INTO `tenants` (`name`) VALUES (?), (?), (?)", dbs[0], dbs[1], dbs[2])
+		require.NoError(t, err)
+
+		cmd := migrateCmd()
+		cmd.AddCommand(migrateApplyCmd())
+		s, err := runCmd(
+			cmd, "apply",
+			"-c", "file://"+path,
+			"--env", "local",
+			"--var", fmt.Sprintf("url=sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "tenants.db")),
+		)
+		require.NoError(t, err)
+		require.Equal(t, 3, strings.Count(s, "Migrating to version 20220318104615 (2 migrations in total)"), "execution per environment")
+		for i := range dbs {
+			_, err = os.Stat(dbs[i])
+			require.NoError(t, err)
+		}
+
+		_, err = db.Exec("INSERT INTO `tenants` (`name`) VALUES (NULL)")
+		require.NoError(t, err)
+		_, err = runCmd(
+			cmd, "apply",
+			"-c", "file://"+path,
+			"--env", "local",
+			"--var", fmt.Sprintf("url=sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "tenants.db")),
+		)
+		// Rows should represent real and consistent values.
+		require.EqualError(t, err, "data.sql.tenants: unsupported row type: <nil>")
+
+		_, err = db.Exec("DELETE FROM `tenants`")
+		require.NoError(t, err)
+		s, err = runCmd(
+			cmd, "apply",
+			"-c", "file://"+path,
+			"--env", "local",
+			"--var", fmt.Sprintf("url=sqlite://file:%s?cache=shared&_fk=1", filepath.Join(p, "tenants.db")),
+		)
+		// Empty list is expanded to zero blocks.
+		require.EqualError(t, err, `env "local" not defined in project file`)
+	})
 }
 
 func TestMigrate_ApplyTxMode(t *testing.T) {
diff --git a/schemahcl/hcl_test.go b/schemahcl/hcl_test.go
index 39315d86..1b8b618b 100644
--- a/schemahcl/hcl_test.go
+++ b/schemahcl/hcl_test.go
@@ -480,7 +480,7 @@ env "dev" {
 `)
 	)
 	require.NoError(t, New(
-		WithDataSource("sql", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, hcl.Diagnostics) {
+		WithDataSource("sql", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
 			attrs, diags := b.Body.JustAttributes()
 			if diags.HasErrors() {
 				return cty.NilVal, diags
@@ -521,7 +521,7 @@ env "dev" {
 func TestDataLocalsRefs(t *testing.T) {
 	var (
 		opts = []Option{
-			WithDataSource("sql", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, hcl.Diagnostics) {
+			WithDataSource("sql", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
 				attrs, diags := b.Body.JustAttributes()
 				if diags.HasErrors() {
 					return cty.NilVal, diags
@@ -532,7 +532,7 @@ func TestDataLocalsRefs(t *testing.T) {
 				}
 				return cty.ObjectVal(map[string]cty.Value{"output": v}), nil
 			}),
-			WithDataSource("text", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, hcl.Diagnostics) {
+			WithDataSource("text", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
 				attrs, diags := b.Body.JustAttributes()
 				if diags.HasErrors() {
 					return cty.NilVal, diags