@@ -94,15 +94,9 @@ func ParseSpec(path string) (PlanSpecs, error) {
9494 return res , err
9595}
9696
97- func generatePlans (specPath string , srcRoot string ) error {
98- specs , err := ParseSpec (specPath )
99- if err != nil {
100- exit (err )
101- }
102- for _ , spec := range specs .Plans {
103- var buf bytes.Buffer
104- fmt .Fprintf (& buf , "// Code generated by plangen.\n \n " )
105- fmt .Fprintf (& buf , `// Copyright 2024 Dolthub, Inc.
97+ func writeHeader (buf * bytes.Buffer , pkg string ) {
98+ _ , _ = fmt .Fprint (buf , "// Code generated by plangen.\n \n " )
99+ _ , _ = fmt .Fprint (buf , `// Copyright 2025 Dolthub, Inc.
106100//
107101// Licensed under the Apache License, Version 2.0 (the "License");
108102// you may not use this file except in compliance with the License.
@@ -115,10 +109,24 @@ func generatePlans(specPath string, srcRoot string) error {
115109// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
116110// See the License for the specific language governing permissions and
117111// limitations under the License.` )
118- fmt .Fprintf (& buf , "\n \n " )
119- fmt .Fprintf (& buf , "package %s\n \n " , * pkg )
112+ _ , _ = fmt .Fprint (buf , "\n \n " )
113+ _ , _ = fmt .Fprintf (buf , "package %s\n \n " , pkg )
114+ }
120115
121- err = generatePlansForSuite (spec , & buf )
116+ func generatePlans (specPath string , srcRoot string ) error {
117+ specs , err := ParseSpec (specPath )
118+ if err != nil {
119+ exit (err )
120+ }
121+ for _ , spec := range specs .Plans {
122+ var buf bytes.Buffer
123+ writeHeader (& buf , * pkg )
124+ if spec .Name == "QueryPlanScriptTests" {
125+ _ , _ = fmt .Fprint (& buf , "import (\n \t \" github.com/dolthub/go-mysql-server/sql\" \n )\n \n " )
126+ err = generatePlansForScriptSuite (spec , & buf )
127+ } else {
128+ err = generatePlansForSuite (spec , & buf )
129+ }
122130 if err != nil {
123131 exit (err )
124132 }
@@ -131,6 +139,33 @@ func generatePlans(specPath string, srcRoot string) error {
131139 return nil
132140}
133141
142+ func writePlanString (w * bytes.Buffer , planString string ) {
143+ for i , line := range strings .Split (planString , "\n " ) {
144+ if i > 0 {
145+ _ , _ = w .WriteString (" + \n " )
146+ }
147+ if len (line ) > 0 {
148+ _ , _ = w .WriteString (fmt .Sprintf (`"%s\n"` , strings .ReplaceAll (line , `"` , `\"` )))
149+ } else {
150+ // final line with comma
151+ _ , _ = w .WriteString ("\" \" ,\n " )
152+ }
153+ }
154+ }
155+
156+ func analyzeQuery (ctx * sql.Context , engine enginetest.QueryEngine , query string ) sql.Node {
157+ binder := planbuilder .New (ctx , engine .EngineAnalyzer ().Catalog , engine .EngineEventScheduler (), nil )
158+ parsed , _ , _ , qFlags , err := binder .Parse (query , nil , false )
159+ if err != nil {
160+ exit (fmt .Errorf ("%w\n failed to parse query: %s" , err , query ))
161+ }
162+ node , err := engine .EngineAnalyzer ().Analyze (ctx , parsed , nil , qFlags )
163+ if err != nil {
164+ exit (fmt .Errorf ("%w\n failed to analyze query: %s" , err , query ))
165+ }
166+ return node
167+ }
168+
134169func generatePlansForSuite (spec PlanSpec , w * bytes.Buffer ) error {
135170 harness := enginetest .NewMemoryHarness ("default" , 1 , 1 , true , nil )
136171 s := specSetup (spec .Name )
@@ -165,49 +200,29 @@ func generatePlansForSuite(spec PlanSpec, w *bytes.Buffer) error {
165200
166201 if ! tt .Skip {
167202 ctx := enginetest .NewContext (harness )
168- binder := planbuilder .New (ctx , engine .EngineAnalyzer ().Catalog , engine .EngineEventScheduler (), nil )
169- parsed , _ , _ , qFlags , err := binder .Parse (tt .Query , nil , false )
170- if err != nil {
171- exit (fmt .Errorf ("%w\n failed to parse query: %s" , err , tt .Query ))
172- }
173- node , err := engine .EngineAnalyzer ().Analyze (ctx , parsed , nil , qFlags )
174- if err != nil {
175- exit (fmt .Errorf ("%w\n failed to analyze query: %s" , err , tt .Query ))
176- }
177-
178- emitPlanString := func (planString string ) {
179- for i , line := range strings .Split (planString , "\n " ) {
180- if i > 0 {
181- _ , _ = w .WriteString (" + \n " )
182- }
183- if len (line ) > 0 {
184- _ , _ = w .WriteString (fmt .Sprintf (`"%s\n"` , strings .ReplaceAll (line , `"` , `\"` )))
185- } else {
186- // final line with comma
187- _ , _ = w .WriteString ("\" \" ,\n " )
188- }
189- }
190- }
191-
203+ node := analyzeQuery (ctx , engine , tt .Query )
192204 _ , _ = w .WriteString (`ExpectedPlan: ` )
193- emitPlanString ( sql .Describe (enginetest .ExtractQueryNode (node ), sql.DescribeOptions {
205+ planString := sql .Describe (enginetest .ExtractQueryNode (node ), sql.DescribeOptions {
194206 Debug : true ,
195- }))
207+ })
208+ writePlanString (w , planString )
196209
197210 if node .IsReadOnly () {
198211 _ , _ = w .WriteString (`ExpectedEstimates: ` )
199- emitPlanString ( sql .Describe (enginetest .ExtractQueryNode (node ), sql.DescribeOptions {
212+ planString = sql .Describe (enginetest .ExtractQueryNode (node ), sql.DescribeOptions {
200213 Estimates : true ,
201- }))
214+ })
215+ writePlanString (w , planString )
202216 err = enginetest .ExecuteNode (ctx , engine , node )
203217 if err != nil {
204218 exit (fmt .Errorf ("%w\n failed to execute query: %s" , err , tt .Query ))
205219 }
206220 _ , _ = w .WriteString (`ExpectedAnalysis: ` )
207- emitPlanString ( sql .Describe (enginetest .ExtractQueryNode (node ), sql.DescribeOptions {
221+ planString = sql .Describe (enginetest .ExtractQueryNode (node ), sql.DescribeOptions {
208222 Analyze : true ,
209223 Estimates : true ,
210- }))
224+ })
225+ writePlanString (w , planString )
211226 }
212227 } else {
213228 _ , _ = w .WriteString (`Skip: true,\n` )
@@ -220,6 +235,71 @@ func generatePlansForSuite(spec PlanSpec, w *bytes.Buffer) error {
220235 return nil
221236}
222237
238+ func generatePlansForScriptSuite (spec PlanSpec , w * bytes.Buffer ) error {
239+ harness := enginetest .NewMemoryHarness ("default" , 1 , 1 , true , nil )
240+ harness .Setup (setup .MydbData )
241+ _ , _ = fmt .Fprintf (w , "var %s = []ScriptTest{\n " , spec .Name )
242+ for _ , tt := range queries .QueryPlanScriptTests {
243+ w .WriteString ("\t {\n " )
244+ if tt .Dialect != "" {
245+ w .WriteString (fmt .Sprintf ("\t \t Dialect: \" %s\" ,\n " , tt .Dialect ))
246+ }
247+ w .WriteString (fmt .Sprintf ("\t \t Name: \" %s\" ,\n " , tt .Name ))
248+ w .WriteString ("\t \t SetUpScript: []string{\n " )
249+ for _ , setupQuery := range tt .SetUpScript {
250+ w .WriteString (fmt .Sprintf ("\t \t \t \" %s\" ,\n " , setupQuery ))
251+ }
252+ w .WriteString ("\t \t },\n " )
253+ w .WriteString ("\t \t Assertions: []ScriptTestAssertion{\n " )
254+ for _ , assertion := range tt .Assertions {
255+ w .WriteString ("\t \t \t {\n " )
256+ if assertion .Skip {
257+ w .WriteString ("\t \t \t \t Skip: true,\n " )
258+ }
259+ w .WriteString (fmt .Sprintf ("\t \t \t \t Query: \" %s\" ,\n " , assertion .Query ))
260+ w .WriteString (fmt .Sprintf ("\t \t \t \t Expected: []sql.Row{\n " ))
261+ for _ , expRow := range assertion .Expected {
262+ w .WriteString (fmt .Sprintf ("\t \t \t \t \t %#v,\n " , expRow ))
263+ }
264+ w .WriteString (fmt .Sprintf ("\t \t \t \t },\n " ))
265+ if assertion .Skip {
266+ w .WriteString ("\t \t \t },\n " )
267+ continue
268+ }
269+
270+ engine , err := harness .NewEngine (nil )
271+ if err != nil {
272+ exit (err )
273+ }
274+ ctx := enginetest .NewContext (harness )
275+ for _ , setupQuery := range tt .SetUpScript {
276+ ctx = ctx .WithQuery (setupQuery )
277+ _ , iter , _ , err := engine .Query (ctx , setupQuery )
278+ if err != nil {
279+ exit (fmt .Errorf ("%w\n failed to execute setup query: %s" , err , setupQuery ))
280+ }
281+ _ , err = sql .RowIterToRows (ctx , iter )
282+ if err != nil {
283+ exit (fmt .Errorf ("%w\n failed to execute setup query: %s" , err , setupQuery ))
284+ }
285+ }
286+
287+ node := analyzeQuery (ctx , engine , assertion .Query )
288+ w .WriteString ("\t \t \t \t ExpectedPlan: " )
289+ planString := sql .Describe (enginetest .ExtractQueryNode (node ), sql.DescribeOptions {
290+ Debug : true ,
291+ })
292+ writePlanString (w , planString )
293+ w .WriteString ("\t \t \t },\n " )
294+ }
295+ w .WriteString ("\t \t },\n " )
296+ w .WriteString ("\t },\n " )
297+ }
298+ w .WriteString ("}" )
299+
300+ return nil
301+ }
302+
223303func specSetup (name string ) [][]setup.SetupScript {
224304 switch name {
225305 case "PlanTests" :
0 commit comments