[no-release-notes] generate query script plans (#3120)

Author: jycor

enginetest/enginetests.go

@@ -793,6 +793,21 @@ func TestQueryPlanWithEngine(t *testing.T, harness Harness, e QueryEngine, tt qu
 	})
 }
 
+func TestQueryPlanScripts(t *testing.T, harness Harness) {
+	harness.Setup(setup.MydbData)
+	for _, script := range queries.QueryPlanScriptTests {
+		if sh, ok := harness.(SkippingHarness); ok {
+			if sh.SkipQueryTest(script.Name) {
+				t.Run(script.Name, func(t *testing.T) {
+					t.Skip(script.Name)
+				})
+				continue
+			}
+		}
+		TestScript(t, harness, script)
+	}
+}
+
 func TestOrderByGroupBy(t *testing.T, harness Harness) {
 	for _, tt := range queries.OrderByGroupByScriptTests {
 		TestScript(t, harness, tt)

enginetest/evaluation.go

@@ -144,12 +144,20 @@ func TestScriptWithEngine(t *testing.T, e QueryEngine, harness Harness, script q
 					}
 					TestQueryWithContext(t, ctx, e, harness, assertion.Query, expected, assertion.ExpectedColumns, assertion.Bindings, nil)
 				}
+
 				if assertion.ExpectedIndexes != nil && !IsServerEngine(e) {
 					evalIndexTest(t, harness, e, assertion.Query, assertion.ExpectedIndexes, assertion.Skip)
 				}
 				if assertion.JoinTypes != nil && !IsServerEngine(e) {
 					evalJoinTypeTest(t, harness, e, assertion.Query, assertion.JoinTypes, assertion.Skip)
 				}
+
+				if assertion.ExpectedPlan != "" && !IsServerEngine(e) {
+					options := sql.DescribeOptions{
+						Debug: true,
+					}
+					TestQueryPlanWithName(t, options.String(), harness, e, assertion.Query, assertion.ExpectedPlan, options)
+				}
 			})
 		}
 	})

enginetest/memory_engine_test.go

@@ -327,6 +327,10 @@ func TestQueryPlans(t *testing.T) {
 	}
 }
 
+func TestQueryPlanScripts(t *testing.T) {
+	enginetest.TestQueryPlanScripts(t, enginetest.NewMemoryHarness("default", 1, testNumPartitions, true, mergableIndexDriver))
+}
+
 func TestSingleQueryPlan(t *testing.T) {
 	t.Skip()
 	tt := []queries.QueryPlanTest{

enginetest/plangen/cmd/plangen/main.go

@@ -94,15 +94,9 @@ func ParseSpec(path string) (PlanSpecs, error) {
 	return res, err
 }
 
-func generatePlans(specPath string, srcRoot string) error {
-	specs, err := ParseSpec(specPath)
-	if err != nil {
-		exit(err)
-	}
-	for _, spec := range specs.Plans {
-		var buf bytes.Buffer
-		fmt.Fprintf(&buf, "// Code generated by plangen.\n\n")
-		fmt.Fprintf(&buf, `// Copyright 2024 Dolthub, Inc.
+func writeHeader(buf *bytes.Buffer, pkg string) {
+	_, _ = fmt.Fprint(buf, "// Code generated by plangen.\n\n")
+	_, _ = fmt.Fprint(buf, `// Copyright 2025 Dolthub, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -115,10 +109,24 @@ func generatePlans(specPath string, srcRoot string) error {
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.`)
-		fmt.Fprintf(&buf, "\n\n")
-		fmt.Fprintf(&buf, "package %s\n\n", *pkg)
+	_, _ = fmt.Fprint(buf, "\n\n")
+	_, _ = fmt.Fprintf(buf, "package %s\n\n", pkg)
+}
 
-		err = generatePlansForSuite(spec, &buf)
+func generatePlans(specPath string, srcRoot string) error {
+	specs, err := ParseSpec(specPath)
+	if err != nil {
+		exit(err)
+	}
+	for _, spec := range specs.Plans {
+		var buf bytes.Buffer
+		writeHeader(&buf, *pkg)
+		if spec.Name == "QueryPlanScriptTests" {
+			_, _ = fmt.Fprint(&buf, "import (\n\t\"github.com/dolthub/go-mysql-server/sql\"\n)\n\n")
+			err = generatePlansForScriptSuite(spec, &buf)
+		} else {
+			err = generatePlansForSuite(spec, &buf)
+		}
 		if err != nil {
 			exit(err)
 		}
@@ -131,6 +139,33 @@ func generatePlans(specPath string, srcRoot string) error {
 	return nil
 }
 
+func writePlanString(w *bytes.Buffer, planString string) {
+	for i, line := range strings.Split(planString, "\n") {
+		if i > 0 {
+			_, _ = w.WriteString(" + \n")
+		}
+		if len(line) > 0 {
+			_, _ = w.WriteString(fmt.Sprintf(`"%s\n"`, strings.ReplaceAll(line, `"`, `\"`)))
+		} else {
+			// final line with comma
+			_, _ = w.WriteString("\"\",\n")
+		}
+	}
+}
+
+func analyzeQuery(ctx *sql.Context, engine enginetest.QueryEngine, query string) sql.Node {
+	binder := planbuilder.New(ctx, engine.EngineAnalyzer().Catalog, engine.EngineEventScheduler(), nil)
+	parsed, _, _, qFlags, err := binder.Parse(query, nil, false)
+	if err != nil {
+		exit(fmt.Errorf("%w\nfailed to parse query: %s", err, query))
+	}
+	node, err := engine.EngineAnalyzer().Analyze(ctx, parsed, nil, qFlags)
+	if err != nil {
+		exit(fmt.Errorf("%w\nfailed to analyze query: %s", err, query))
+	}
+	return node
+}
+
 func generatePlansForSuite(spec PlanSpec, w *bytes.Buffer) error {
 	harness := enginetest.NewMemoryHarness("default", 1, 1, true, nil)
 	s := specSetup(spec.Name)
@@ -165,49 +200,29 @@ func generatePlansForSuite(spec PlanSpec, w *bytes.Buffer) error {
 
 		if !tt.Skip {
 			ctx := enginetest.NewContext(harness)
-			binder := planbuilder.New(ctx, engine.EngineAnalyzer().Catalog, engine.EngineEventScheduler(), nil)
-			parsed, _, _, qFlags, err := binder.Parse(tt.Query, nil, false)
-			if err != nil {
-				exit(fmt.Errorf("%w\nfailed to parse query: %s", err, tt.Query))
-			}
-			node, err := engine.EngineAnalyzer().Analyze(ctx, parsed, nil, qFlags)
-			if err != nil {
-				exit(fmt.Errorf("%w\nfailed to analyze query: %s", err, tt.Query))
-			}
-
-			emitPlanString := func(planString string) {
-				for i, line := range strings.Split(planString, "\n") {
-					if i > 0 {
-						_, _ = w.WriteString(" + \n")
-					}
-					if len(line) > 0 {
-						_, _ = w.WriteString(fmt.Sprintf(`"%s\n"`, strings.ReplaceAll(line, `"`, `\"`)))
-					} else {
-						// final line with comma
-						_, _ = w.WriteString("\"\",\n")
-					}
-				}
-			}
-
+			node := analyzeQuery(ctx, engine, tt.Query)
 			_, _ = w.WriteString(`ExpectedPlan: `)
-			emitPlanString(sql.Describe(enginetest.ExtractQueryNode(node), sql.DescribeOptions{
+			planString := sql.Describe(enginetest.ExtractQueryNode(node), sql.DescribeOptions{
 				Debug: true,
-			}))
+			})
+			writePlanString(w, planString)
 
 			if node.IsReadOnly() {
 				_, _ = w.WriteString(`ExpectedEstimates: `)
-				emitPlanString(sql.Describe(enginetest.ExtractQueryNode(node), sql.DescribeOptions{
+				planString = sql.Describe(enginetest.ExtractQueryNode(node), sql.DescribeOptions{
 					Estimates: true,
-				}))
+				})
+				writePlanString(w, planString)
 				err = enginetest.ExecuteNode(ctx, engine, node)
 				if err != nil {
 					exit(fmt.Errorf("%w\nfailed to execute query: %s", err, tt.Query))
 				}
 				_, _ = w.WriteString(`ExpectedAnalysis: `)
-				emitPlanString(sql.Describe(enginetest.ExtractQueryNode(node), sql.DescribeOptions{
+				planString = sql.Describe(enginetest.ExtractQueryNode(node), sql.DescribeOptions{
 					Analyze:   true,
 					Estimates: true,
-				}))
+				})
+				writePlanString(w, planString)
 			}
 		} else {
 			_, _ = w.WriteString(`Skip: true,\n`)
@@ -220,6 +235,71 @@ func generatePlansForSuite(spec PlanSpec, w *bytes.Buffer) error {
 	return nil
 }
 
+func generatePlansForScriptSuite(spec PlanSpec, w *bytes.Buffer) error {
+	harness := enginetest.NewMemoryHarness("default", 1, 1, true, nil)
+	harness.Setup(setup.MydbData)
+	_, _ = fmt.Fprintf(w, "var %s = []ScriptTest{\n", spec.Name)
+	for _, tt := range queries.QueryPlanScriptTests {
+		w.WriteString("\t{\n")
+		if tt.Dialect != "" {
+			w.WriteString(fmt.Sprintf("\t\tDialect: \"%s\",\n", tt.Dialect))
+		}
+		w.WriteString(fmt.Sprintf("\t\tName: \"%s\",\n", tt.Name))
+		w.WriteString("\t\tSetUpScript: []string{\n")
+		for _, setupQuery := range tt.SetUpScript {
+			w.WriteString(fmt.Sprintf("\t\t\t\"%s\",\n", setupQuery))
+		}
+		w.WriteString("\t\t},\n")
+		w.WriteString("\t\tAssertions: []ScriptTestAssertion{\n")
+		for _, assertion := range tt.Assertions {
+			w.WriteString("\t\t\t{\n")
+			if assertion.Skip {
+				w.WriteString("\t\t\t\tSkip: true,\n")
+			}
+			w.WriteString(fmt.Sprintf("\t\t\t\tQuery: \"%s\",\n", assertion.Query))
+			w.WriteString(fmt.Sprintf("\t\t\t\tExpected: []sql.Row{\n"))
+			for _, expRow := range assertion.Expected {
+				w.WriteString(fmt.Sprintf("\t\t\t\t\t%#v,\n", expRow))
+			}
+			w.WriteString(fmt.Sprintf("\t\t\t\t},\n"))
+			if assertion.Skip {
+				w.WriteString("\t\t\t},\n")
+				continue
+			}
+
+			engine, err := harness.NewEngine(nil)
+			if err != nil {
+				exit(err)
+			}
+			ctx := enginetest.NewContext(harness)
+			for _, setupQuery := range tt.SetUpScript {
+				ctx = ctx.WithQuery(setupQuery)
+				_, iter, _, err := engine.Query(ctx, setupQuery)
+				if err != nil {
+					exit(fmt.Errorf("%w\nfailed to execute setup query: %s", err, setupQuery))
+				}
+				_, err = sql.RowIterToRows(ctx, iter)
+				if err != nil {
+					exit(fmt.Errorf("%w\nfailed to execute setup query: %s", err, setupQuery))
+				}
+			}
+
+			node := analyzeQuery(ctx, engine, assertion.Query)
+			w.WriteString("\t\t\t\tExpectedPlan: ")
+			planString := sql.Describe(enginetest.ExtractQueryNode(node), sql.DescribeOptions{
+				Debug: true,
+			})
+			writePlanString(w, planString)
+			w.WriteString("\t\t\t},\n")
+		}
+		w.WriteString("\t\t},\n")
+		w.WriteString("\t},\n")
+	}
+	w.WriteString("}")
+
+	return nil
+}
+
 func specSetup(name string) [][]setup.SetupScript {
 	switch name {
 	case "PlanTests":

enginetest/plangen/testdata/spec.yaml

@@ -16,4 +16,6 @@ plans:
   - name: GeneratedColumnPlanTests
     path: enginetest/queries/generated_column_plans.go
   - name: SysbenchPlanTests
-    path: enginetest/queries/sysbench_plans.go
+    path: enginetest/queries/sysbench_plans.go
+  - name: QueryPlanScriptTests
+    path: enginetest/queries/query_plan_script_tests.go