Repository: ariga/atlas
Pull request #1180 (2022-10-08 16:03:42)

schemahcl: add support for datasource and locals


Patch:
diff --git a/doc/md/guides/postgres/included-columns.md b/doc/md/guides/postgres/included-columns.md
index 765e31a0..8e21863d 100644
--- a/doc/md/guides/postgres/included-columns.md
+++ b/doc/md/guides/postgres/included-columns.md
@@ -4,7 +4,9 @@ title: Indexes with Included Columns in PostgreSQL
 slug: /guides/postgres/included-columns
 ---
 
-With PostgreSQL, we can create indexes with `INCLUDE` clauses, which are types of indexes that specify a list of columns to be included in the index as non-key columns. If used correctly, indexes with included columns improve performance and reduce total costs.
+With PostgreSQL, we can create *covering indexes* using the `INCLUDE` clause, which are types of indexes that specify a
+list of columns to be included in the index as non-key columns. If used correctly, indexes with included columns improve
+performance and reduce total costs.
 
 ### Basic PostgreSQL syntax for using `INCLUDE` clause with an index:
 
@@ -16,14 +18,17 @@ INCLUDE(included_column_list);
 
 ### How do they work?
 
-A B-Tree index is made of two structures:
-1. The B-tree
-2. The doubly linked list at the leaf node level of the B-tree
+In PostgreSQL, a B-Tree index creates a multi-level tree structure where each level can be used as a doubly-linked list
+of pages. Leaf pages are those at the lowest level of a tree, that point to rows of tables.
 
-In an index with included columns, records of the columns mentioned in the `INCLUDE` clause are included in the doubly linked list at the leaf node level of the B-tree, corresponding to the heap (rows/tuples stored in a table).
+With *covering indexes*, records of the columns mentioned in the `INCLUDE` clause are included in the leaf pages of the
+B-Tree as "payload" and are not part of the search key.
 
 :::info
-Each index is stored separately from the table's main data area, which in PostgreSQL this is known as the table's _heap_.  Leaf pages are the pages on the lowest level of the index tree. To learn more about the B-tree index structure in PostgreSQL, visit the PostgreSQL documentation [here](https://www.postgresql.org/docs/current/btree-implementation.html#BTREE-STRUCTURE). 
+Each index is stored separately from the table's main data area, which in PostgreSQL this is known as the table's _heap_.
+To learn more about the PostgreSQL B-tree index structure and *covering indexes*, visit the documentation:
+1. [B-Tree structure](https://www.postgresql.org/docs/current/btree-implementation.html#BTREE-STRUCTURE).
+2. [Covering indexes](https://www.postgresql.org/docs/current/indexes-index-only-scans.html).
 :::
 
 ### When do we need them?
diff --git a/doc/website/sidebars.js b/doc/website/sidebars.js
index efc067e3..6edbe9f3 100644
--- a/doc/website/sidebars.js
+++ b/doc/website/sidebars.js
@@ -125,7 +125,7 @@ module.exports = {
                 {
                     type: 'doc',
                     id: 'guides/postgres/included-columns',
-                    label: 'Included Columns'
+                    label: 'Covering Indexes'
                 },
             ],
         },
diff --git a/schemahcl/context.go b/schemahcl/context.go
index d23d847c..bc8f9dbe 100644
--- a/schemahcl/context.go
+++ b/schemahcl/context.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"reflect"
 	"strconv"
+	"strings"
 
 	"github.com/hashicorp/hcl/v2"
 	"github.com/hashicorp/hcl/v2/gohcl"
@@ -16,8 +17,8 @@ import (
 	"github.com/zclconf/go-cty/cty/convert"
 )
 
-// varDef is an HCL resource that defines an input variable to the Atlas DDL document.
-type varDef struct {
+// blockVar is an HCL resource that defines an input variable to the Atlas DDL document.
+type blockVar struct {
 	Name    string    `hcl:",label"`
 	Type    cty.Value `hcl:"type"`
 	Default cty.Value `hcl:"default,optional"`
@@ -27,14 +28,14 @@ type varDef struct {
 // input variables in the document body by defining "variable" blocks:
 //
 //	variable "name" {
-//	  type = string // also supported: int, bool
+//	  type = string // also supported: number, bool
 //	  default = "rotemtam"
 //	}
 func (s *State) setInputVals(ctx *hcl.EvalContext, body hcl.Body, input map[string]cty.Value) error {
 	var (
 		doc struct {
-			Vars   []*varDef `hcl:"variable,block"`
-			Remain hcl.Body  `hcl:",remain"`
+			Vars   []*blockVar `hcl:"variable,block"`
+			Remain hcl.Body    `hcl:",remain"`
 		}
 		nctx = varBlockContext(ctx)
 	)
@@ -68,16 +69,130 @@ func (s *State) setInputVals(ctx *hcl.EvalContext, body hcl.Body, input map[stri
 	return nil
 }
 
+// evalReferences evaluates data blocks.
+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)
+	}
+	var (
+		nodes  = make(map[[3]string]*node)
+		blocks = make(hclsyntax.Blocks, 0, len(body.Blocks))
+	)
+	for _, b := range body.Blocks {
+		switch b := b; b.Type {
+		case dataBlock:
+			if len(b.Labels) < 2 {
+				return fmt.Errorf("data block %q must have exactly 2 labels", b.Type)
+			}
+			h, ok := s.config.datasrc[b.Labels[0]]
+			if !ok {
+				return fmt.Errorf("missing data source handler for %q", b.Labels[0])
+			}
+			// Data references are combined from
+			// "data", "source" and "name" labels.
+			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) },
+				edges: func() []hcl.Traversal { return bodyVars(b.Body) },
+			}
+		case localsBlock:
+			for k, v := range b.Body.Attributes {
+				k, v := k, v
+				// Local references are combined from
+				// "local" and "name" labels.
+				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) },
+				}
+			}
+		default:
+			blocks = append(blocks, b)
+		}
+	}
+	var (
+		visit    func(n *node) error
+		visited  = make(map[*node]bool)
+		progress = make(map[*node]bool)
+	)
+	visit = func(n *node) error {
+		if visited[n] {
+			return nil
+		}
+		if progress[n] {
+			addr := n.addr[:]
+			if addr[2] == "" {
+				addr = addr[:2]
+			}
+			return fmt.Errorf("cyclic reference to %q", strings.Join(addr, "."))
+		}
+		progress[n] = true
+		for _, e := range n.edges() {
+			var addr [3]string
+			switch root := e.RootName(); {
+			case root == localRef && len(e) == 2:
+				addr = [3]string{localRef, e[1].(hcl.TraverseAttr).Name, ""}
+			case root == dataBlock && len(e) > 2:
+				addr = [3]string{dataBlock, e[1].(hcl.TraverseAttr).Name, e[2].(hcl.TraverseAttr).Name}
+			}
+			// Unrecognized reference.
+			if nodes[addr] == nil {
+				continue
+			}
+			if err := visit(nodes[addr]); err != nil {
+				return err
+			}
+		}
+		delete(progress, n)
+		v, diags := n.value()
+		if diags.HasErrors() {
+			return diags
+		}
+		switch n.addr[0] {
+		case dataBlock:
+			data := make(map[string]cty.Value)
+			if vv, ok := ctx.Variables[dataBlock]; ok {
+				data = vv.AsValueMap()
+			}
+			src := make(map[string]cty.Value)
+			if vv, ok := data[n.addr[1]]; ok {
+				src = vv.AsValueMap()
+			}
+			src[n.addr[2]] = v
+			data[n.addr[1]] = cty.ObjectVal(src)
+			ctx.Variables[dataBlock] = cty.ObjectVal(data)
+		case localRef:
+			locals := make(map[string]cty.Value)
+			if vv, ok := ctx.Variables[localRef]; ok {
+				locals = vv.AsValueMap()
+			}
+			locals[n.addr[1]] = v
+			ctx.Variables[localRef] = cty.ObjectVal(locals)
+		}
+		return nil
+	}
+	for _, n := range nodes {
+		if err := visit(n); err != nil {
+			return err
+		}
+	}
+	body.Blocks = blocks
+	return nil
+}
+
 func mergeCtxVar(ctx *hcl.EvalContext, vals map[string]cty.Value) {
-	const key = "var"
-	v, ok := ctx.Variables[key]
+	v, ok := ctx.Variables[varRef]
 	if ok {
 		v.ForEachElement(func(key cty.Value, val cty.Value) (stop bool) {
 			vals[key.AsString()] = val
 			return false
 		})
 	}
-	ctx.Variables[key] = cty.ObjectVal(vals)
+	ctx.Variables[varRef] = cty.ObjectVal(vals)
 }
 
 func setBlockVars(ctx *hcl.EvalContext, b *hclsyntax.Body) (*hcl.EvalContext, error) {
@@ -229,6 +344,10 @@ var (
 // Built-in blocks.
 const (
 	varBlock     = "variable"
+	dataBlock    = "data"
+	localsBlock  = "locals"
+	varRef       = "var"
+	localRef     = "local"
 	dynamicBlock = "dynamic"
 	forEachAttr  = "for_each"
 )
@@ -302,3 +421,13 @@ func extractDef(blk *hclsyntax.Block, parent *blockDef) *blockDef {
 	}
 	return cur
 }
+
+func bodyVars(b *hclsyntax.Body) (vars []hcl.Traversal) {
+	for _, attr := range b.Attributes {
+		vars = append(vars, hclsyntax.Variables(attr.Expr)...)
+	}
+	for _, b := range b.Blocks {
+		vars = append(vars, bodyVars(b.Body)...)
+	}
+	return
+}
diff --git a/schemahcl/hcl.go b/schemahcl/hcl.go
index f532c0a6..7a3b8627 100644
--- a/schemahcl/hcl.go
+++ b/schemahcl/hcl.go
@@ -85,13 +85,16 @@ func (s *State) Eval(parsed *hclparse.Parser, v any, input map[string]cty.Value)
 			return err
 		}
 		body := file.Body.(*hclsyntax.Body)
+		if err := s.evalReferences(ctx, body); err != nil {
+			return err
+		}
 		for _, blk := range body.Blocks {
-			// Variable definition blocks are available in the HCL source but not reachable by reference.
-			if blk.Type == varBlock {
-				continue
+			// Variable blocks are available in the HCL
+			// source but not reachable by reference.
+			if blk.Type != varBlock {
+				allBlocks = append(allBlocks, blk)
+				reg.child(extractDef(blk, reg))
 			}
-			allBlocks = append(allBlocks, blk)
-			reg.child(extractDef(blk, reg))
 		}
 	}
 	vars, err := blockVars(allBlocks, "", reg)
diff --git a/schemahcl/opts.go b/schemahcl/opts.go
index 73fc1c0f..8663bf51 100644
--- a/schemahcl/opts.go
+++ b/schemahcl/opts.go
@@ -10,6 +10,7 @@ import (
 	"reflect"
 
 	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/hcl/v2/hclsyntax"
 	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/function"
 )
@@ -20,6 +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)
 	}
 
 	// Option configures a Config.
@@ -67,6 +69,34 @@ func WithScopedEnums(path string, enums ...string) Option {
 	}
 }
 
+// WithDataSource registers a data source name and its corresponding handler.
+// e.g., the example below registers a data source named "text" that returns
+// the string defined in the data source block.
+//
+//	WithDataSource("text", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, hcl.Diagnostics) {
+//		attrs, diags := b.Body.JustAttributes()
+//		if diags.HasErrors() {
+//			return cty.NilVal, diags
+//		}
+//		v, diags := attrs["value"].Expr.Value(ctx)
+//		if diags.HasErrors() {
+//			return cty.NilVal, diags
+//		}
+//		return cty.ObjectVal(map[string]cty.Value{"output": v}), nil
+//	}),
+//
+//	data "text" "hello" {
+//	  value = "hello world"
+//	}
+func WithDataSource(name string, h func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, hcl.Diagnostics)) 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[name] = h
+	}
+}
+
 // WithTypes configures the list of given types as identifiers in the unmarshaling context.
 func WithTypes(typeSpecs []*TypeSpec) Option {
 	newCtx := func() *hcl.EvalContext {


Test patch:
diff --git a/schemahcl/hcl_test.go b/schemahcl/hcl_test.go
index 2b9d7fff..39315d86 100644
--- a/schemahcl/hcl_test.go
+++ b/schemahcl/hcl_test.go
@@ -11,6 +11,8 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/hcl/v2/hclsyntax"
 	"github.com/stretchr/testify/require"
 	"github.com/zclconf/go-cty/cty"
 )
@@ -477,7 +479,19 @@ env "dev" {
 }
 `)
 	)
-	require.NoError(t, New().EvalBytes(b, &doc, nil))
+	require.NoError(t, New(
+		WithDataSource("sql", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, hcl.Diagnostics) {
+			attrs, diags := b.Body.JustAttributes()
+			if diags.HasErrors() {
+				return cty.NilVal, diags
+			}
+			v, diags := attrs["query"].Expr.Value(ctx)
+			if diags.HasErrors() {
+				return cty.NilVal, diags
+			}
+			return cty.ObjectVal(map[string]cty.Value{"query": v}), nil
+		}),
+	).EvalBytes(b, &doc, nil))
 	require.Len(t, doc.Envs, 6)
 	require.Equal(t, "prod", doc.Envs[0].Name)
 	require.EqualValues(t, doc.Envs[0].URL, "mysql://root:pass@:3306/atlas")
@@ -503,3 +517,113 @@ env "dev" {
 	})
 	require.EqualError(t, err, `variable "domains": a number is required`)
 }
+
+func TestDataLocalsRefs(t *testing.T) {
+	var (
+		opts = []Option{
+			WithDataSource("sql", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, hcl.Diagnostics) {
+				attrs, diags := b.Body.JustAttributes()
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				v, diags := attrs["result"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				return cty.ObjectVal(map[string]cty.Value{"output": v}), nil
+			}),
+			WithDataSource("text", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, hcl.Diagnostics) {
+				attrs, diags := b.Body.JustAttributes()
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				v, diags := attrs["value"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				return cty.ObjectVal(map[string]cty.Value{"output": v}), nil
+			}),
+		}
+		doc struct {
+			Values []string `spec:"vs"`
+		}
+		b = []byte(`
+variable "url" {
+  type    = string
+  default = "mysql://root:pass@:3306/atlas"
+}
+
+locals {
+  a = "local-a"
+  // locals can reference other locals.
+  b = "local-b-ref-local-a: ${local.a}"
+  // locals can reference data sources.
+  c = "local-c-ref-data-a: ${data.text.a.output}"
+  d = "local-d"
+}
+
+data "sql" "tenants" {
+  url = var.url
+  // language=mysql
+  query = <<EOS
+SELECT schema_name
+  FROM information_schema.schemata
+  WHERE schema_name LIKE 'tenant_%'
+EOS
+  // fake result.
+  result = "data-sql-tenants"
+}
+
+data "text" "a" {
+  // data sources can reference data sources.
+  value = "data-text-a-ref-data-sql-tenants: ${data.sql.tenants.output}"
+}
+
+data "text" "b" {
+  // data sources can reference locals.
+  value = "data-text-b-ref-local-d: ${local.d}"
+}
+
+vs = [
+  local.a,
+  local.b,
+  local.c,
+  data.sql.tenants.output,
+  data.text.a.output,
+  data.text.b.output,
+]
+`)
+	)
+	require.NoError(t, New(opts...).EvalBytes(b, &doc, nil))
+	require.Equal(t, []string{
+		"local-a",
+		"local-b-ref-local-a: local-a",
+		"local-c-ref-data-a: data-text-a-ref-data-sql-tenants: data-sql-tenants",
+		"data-sql-tenants",
+		"data-text-a-ref-data-sql-tenants: data-sql-tenants",
+		"data-text-b-ref-local-d: local-d",
+	}, doc.Values)
+
+	b = []byte(`locals { a = local.a }`)
+	require.EqualError(t, New(opts...).EvalBytes(b, &doc, nil), `cyclic reference to "local.a"`)
+
+	b = []byte(`
+locals {
+  a = "a"
+  b = local.c
+  c = local.b
+}
+`)
+	require.Error(t, New(opts...).EvalBytes(b, &doc, nil), `cyclic reference to "local.c"`)
+
+	b = []byte(`
+data "text" "a" {
+  value = local.a
+}
+
+locals {
+  a = data.text.a.output
+}
+`)
+	require.Error(t, New(opts...).EvalBytes(b, &doc, nil), `cyclic reference to "data.text.a"`)
+}