Added InjectedStatement as an AST node

dolthub · May 20, 2024 · de5a32a · de5a32a
1 parent 47f2ce2
commit de5a32a
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 3 deletions.
diff --git a/go.mod b/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/dolthub/go-icu-regex v0.0.0-20230524105445-af7e7991c97e
 	github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71
 	github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81
-	github.com/dolthub/vitess v0.0.0-20240514182535-00461f67afae
+	github.com/dolthub/vitess v0.0.0-20240520121817-ae2979fa661b
 	github.com/go-kit/kit v0.10.0
 	github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d
 	github.com/gocraft/dbr/v2 v2.7.2

diff --git a/go.sum b/go.sum
@@ -58,8 +58,8 @@ github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71 h1:bMGS25NWAGTE
 github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71/go.mod h1:2/2zjLQ/JOOSbbSboojeg+cAwcRV0fDLzIiWch/lhqI=
 github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81 h1:7/v8q9XGFa6q5Ap4Z/OhNkAMBaK5YeuEzwJt+NZdhiE=
 github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81/go.mod h1:siLfyv2c92W1eN/R4QqG/+RjjX5W2+gCTRjZxBjI3TY=
-github.com/dolthub/vitess v0.0.0-20240514182535-00461f67afae h1:q3Fy11eAibKoO6t2c+GADn1qafVLPSwGE349DVSTngI=
-github.com/dolthub/vitess v0.0.0-20240514182535-00461f67afae/go.mod h1:uBvlRluuL+SbEWTCZ68o0xvsdYZER3CEG/35INdzfJM=
+github.com/dolthub/vitess v0.0.0-20240520121817-ae2979fa661b h1:Vp41KCx2+1n4ftEfadtDLDo9aKKFel+YviSyrIc7yy4=
+github.com/dolthub/vitess v0.0.0-20240520121817-ae2979fa661b/go.mod h1:uBvlRluuL+SbEWTCZ68o0xvsdYZER3CEG/35INdzfJM=
 github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/eapache/go-resiliency v1.1.0/go.mod h1:kFI+JgMyC7bLPUVY133qvEBtVayf5mFgVsvEsIPBvNs=
 github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21/go.mod h1:+020luEh2TKB4/GOp8oxxtq0Daoen/Cii55CzbTV6DU=

diff --git a/sql/planbuilder/builder.go b/sql/planbuilder/builder.go
@@ -15,6 +15,7 @@
 package planbuilder
 
 import (
+	"fmt"
 	"strings"
 	"sync"
 
@@ -361,6 +362,8 @@ func (b *Builder) build(inScope *scope, stmt ast.Statement, query string) (outSc
 		return b.buildExecute(inScope, n)
 	case *ast.Deallocate:
 		return b.buildDeallocate(inScope, n)
+	case ast.InjectedStatement:
+		return b.buildInjectedStatement(inScope, n)
 	}
 	return
 }
@@ -400,6 +403,26 @@ func (b *Builder) buildVirtualTableScan(db string, tab sql.Table) *plan.VirtualC
 	return plan.NewVirtualColumnTable(tab, projections)
 }
 
+// buildInjectedStatement returns the sql.Node encapsulated by the injected statement.
+func (b *Builder) buildInjectedStatement(inScope *scope, n ast.InjectedStatement) (outScope *scope) {
+	resolvedChildren := make([]any, len(n.Children))
+	for i, child := range n.Children {
+		resolvedChildren[i] = b.buildScalar(inScope, child)
+	}
+	stmt, err := n.Statement.WithResolvedChildren(resolvedChildren)
+	if err != nil {
+		b.handleErr(err)
+		return nil
+	}
+	if sqlNode, ok := stmt.(sql.ExecSourceRel); ok {
+		outScope = inScope.push()
+		outScope.node = sqlNode
+		return outScope
+	}
+	b.handleErr(fmt.Errorf("Injected statement does not resolve to a valid node"))
+	return nil
+}
+
 // assignColumnIndexes fixes the column indexes in the expression to match the schema given
 func assignColumnIndexes(e sql.Expression, schema sql.Schema) sql.Expression {
 	e, _, _ = transform.Expr(e, func(e sql.Expression) (sql.Expression, transform.TreeIdentity, error) {

diff --git a/sql/planbuilder/scalar.go b/sql/planbuilder/scalar.go
@@ -250,6 +250,7 @@ func (b *Builder) buildScalar(inScope *scope, e ast.Expr) (ex sql.Expression) {
 		expr, err := v.Expression.WithResolvedChildren(resolvedChildren)
 		if err != nil {
 			b.handleErr(err)
+			return nil
 		}
 		if sqlExpr, ok := expr.(sql.Expression); ok {
 			return sqlExpr