Fix prep AS OF (#1054)

Author: max-hoffman

engine.go

@@ -321,18 +321,18 @@ func (e *Engine) analyzeQuery(ctx *sql.Context, query string, parsed sql.Node, b
 		return nil, err
 	}
 
-	analyzed, err = e.Analyzer.Analyze(ctx, parsed, nil)
-	if err != nil {
-		return nil, err
-	}
-
 	if len(bindings) > 0 {
-		analyzed, err = plan.ApplyBindings(analyzed, bindings)
+		parsed, err = plan.ApplyBindings(parsed, bindings)
 		if err != nil {
 			return nil, err
 		}
 	}
 
+	analyzed, err = e.Analyzer.Analyze(ctx, parsed, nil)
+	if err != nil {
+		return nil, err
+	}
+
 	return analyzed, nil
 }
 

enginetest/evaluation.go

@@ -256,6 +256,10 @@ func TestQueryWithEngine(t *testing.T, harness Harness, e *sqle.Engine, tt queri
 func TestQueryWithContext(t *testing.T, ctx *sql.Context, e *sqle.Engine, q string, expected []sql.Row, expectedCols []*sql.Column, bindings map[string]sql.Expression) {
 	ctx = ctx.WithQuery(q)
 	require := require.New(t)
+	if len(bindings) > 0 {
+		_, err := e.PrepareQuery(ctx, q)
+		require.NoError(err)
+	}
 	sch, iter, err := e.QueryWithBindings(ctx, q, bindings)
 	require.NoError(err, "Unexpected error for query %s", q)
 

enginetest/initialization.go

@@ -177,7 +177,7 @@ func RunEngineScripts(ctx *sql.Context, e *sqle.Engine, scripts []setup.SetupScr
 	return e, nil
 }
 
-func MustQuery(ctx *sql.Context, e *sqle.Engine, q string) []sql.Row {
+func MustQuery(ctx *sql.Context, e *sqle.Engine, q string) (sql.Schema, []sql.Row) {
 	sch, iter, err := e.Query(ctx, q)
 	if err != nil {
 		panic(err)
@@ -186,7 +186,42 @@ func MustQuery(ctx *sql.Context, e *sqle.Engine, q string) []sql.Row {
 	if err != nil {
 		panic(err)
 	}
-	return rows
+	return sch, rows
+}
+
+func MustQueryWithBindings(ctx *sql.Context, e *sqle.Engine, q string, bindings map[string]sql.Expression) (sql.Schema, []sql.Row) {
+	ctx = ctx.WithQuery(q)
+	sch, iter, err := e.QueryWithBindings(ctx, q, bindings)
+	if err != nil {
+		panic(err)
+	}
+
+	rows, err := sql.RowIterToRows(ctx, sch, iter)
+	if err != nil {
+		panic(err)
+	}
+
+	return sch, rows
+}
+
+func MustQueryWithPreBindings(ctx *sql.Context, e *sqle.Engine, q string, bindings map[string]sql.Expression) (sql.Node, sql.Schema, []sql.Row) {
+	ctx = ctx.WithQuery(q)
+	pre, err := e.PrepareQuery(ctx, q)
+	if err != nil {
+		panic(err)
+	}
+
+	sch, iter, err := e.QueryWithBindings(ctx, q, bindings)
+	if err != nil {
+		panic(err)
+	}
+
+	rows, err := sql.RowIterToRows(ctx, sch, iter)
+	if err != nil {
+		panic(err)
+	}
+
+	return pre, sch, rows
 }
 
 func mustNewEngine(t *testing.T, h Harness) *sqle.Engine {

enginetest/testgen_test.go

@@ -175,8 +175,8 @@ func TestWriteCreateTableQueries(t *testing.T) {
 	for _, tt := range queries.CreateTableQueries {
 		ctx := NewContext(harness)
 		engine := mustNewEngine(t, harness)
-		_ = MustQuery(ctx, engine, tt.WriteQuery)
-		res := MustQuery(ctx, engine, tt.SelectQuery)
+		_, _ = MustQuery(ctx, engine, tt.WriteQuery)
+		_, res := MustQuery(ctx, engine, tt.SelectQuery)
 
 		w.WriteString("  {\n")
 		w.WriteString(fmt.Sprintf("    WriteQuery:`%s`,\n", tt.WriteQuery))

sql/analyzer/resolve_tables.go

@@ -268,7 +268,11 @@ func reresolveTables(ctx *sql.Context, a *Analyzer, node sql.Node, scope *Scope,
 			if n.Database != nil {
 				db = n.Database.Name()
 			}
-			to, err = resolveTable(ctx, plan.NewUnresolvedTable(n.Name(), db), a)
+			var asof sql.Expression
+			if n.AsOf != nil {
+				asof = expression.NewLiteral(n.AsOf, nil)
+			}
+			to, err = resolveTable(ctx, plan.NewUnresolvedTableAsOf(n.Name(), db, asof), a)
 			if err != nil {
 				return nil, transform.SameTree, err
 			}