Merge branch 'main' into zachmu/enginetests5

Author: zachmu

engine.go

@@ -209,7 +209,7 @@ func (e *Engine) AnalyzeQuery(
 	ctx *sql.Context,
 	query string,
 ) (sql.Node, error) {
-	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.Parser)
+	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.EventScheduler, e.Parser)
 	parsed, _, _, qFlags, err := binder.Parse(query, nil, false)
 	if err != nil {
 		return nil, err
@@ -237,7 +237,7 @@ func (e *Engine) PrepareParsedQuery(
 	statementKey, query string,
 	stmt sqlparser.Statement,
 ) (sql.Node, error) {
-	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.Parser)
+	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.EventScheduler, e.Parser)
 	node, _, err := binder.BindOnly(stmt, query, nil)
 
 	if err != nil {
@@ -443,13 +443,24 @@ func (e *Engine) QueryWithBindings(ctx *sql.Context, query string, parsed sqlpar
 		if err2 != nil {
 			return nil, nil, nil, errors.Wrap(err, "unable to clear autocommit transaction: "+err2.Error())
 		}
+		return nil, nil, nil, err
+	}
 
+	var schema sql.Schema
+	iter, schema = rowexec.FinalizeIters(ctx, analyzed, qFlags, iter)
+	if err != nil {
+		clearAutocommitErr := clearAutocommitTransaction(ctx)
+		if clearAutocommitErr != nil {
+			return nil, nil, nil, errors.Wrap(err, "unable to clear autocommit transaction: "+clearAutocommitErr.Error())
+		}
 		return nil, nil, nil, err
 	}
 
-	iter = finalizeIters(ctx, analyzed, qFlags, iter)
+	if schema == nil {
+		schema = analyzed.Schema()
+	}
 
-	return analyzed.Schema(), iter, qFlags, nil
+	return schema, iter, qFlags, nil
 }
 
 // PrepQueryPlanForExecution prepares a query plan for execution and returns the result schema with a row iterator to
@@ -478,13 +489,24 @@ func (e *Engine) PrepQueryPlanForExecution(ctx *sql.Context, _ string, plan sql.
 		if err2 != nil {
 			return nil, nil, nil, errors.Wrap(err, "unable to clear autocommit transaction: "+err2.Error())
 		}
+		return nil, nil, nil, err
+	}
 
+	var schema sql.Schema
+	iter, schema = rowexec.FinalizeIters(ctx, plan, qFlags, iter)
+	if err != nil {
+		clearAutocommitErr := clearAutocommitTransaction(ctx)
+		if clearAutocommitErr != nil {
+			return nil, nil, nil, errors.Wrap(err, "unable to clear autocommit transaction: "+clearAutocommitErr.Error())
+		}
 		return nil, nil, nil, err
 	}
 
-	iter = finalizeIters(ctx, plan, qFlags, iter)
+	if schema == nil {
+		schema = plan.Schema()
+	}
 
-	return plan.Schema(), iter, qFlags, nil
+	return schema, iter, qFlags, nil
 }
 
 // BoundQueryPlan returns query plan for the given statement with the given bindings applied
@@ -495,7 +517,7 @@ func (e *Engine) BoundQueryPlan(ctx *sql.Context, query string, parsed sqlparser
 
 	query = sql.RemoveSpaceAndDelimiter(query, ';')
 
-	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.Parser)
+	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.EventScheduler, e.Parser)
 	binder.SetBindings(bindings)
 
 	// Begin a transaction if necessary (no-op if one is in flight)
@@ -549,7 +571,7 @@ func (e *Engine) preparedStatement(ctx *sql.Context, query string, parsed sqlpar
 		preparedAst, preparedDataFound = e.PreparedDataCache.GetCachedStmt(ctx.Session.ID(), query)
 	}
 
-	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.Parser)
+	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.EventScheduler, e.Parser)
 	if preparedDataFound {
 		parsed = preparedAst
 		binder.SetBindings(bindings)
@@ -725,18 +747,19 @@ func (e *Engine) CloseSession(connID uint32) {
 }
 
 func (e *Engine) beginTransaction(ctx *sql.Context) error {
-	beginNewTransaction := ctx.GetTransaction() == nil || plan.ReadCommitted(ctx)
-	if beginNewTransaction {
-		ctx.GetLogger().Tracef("beginning new transaction")
-		ts, ok := ctx.Session.(sql.TransactionSession)
-		if ok {
-			tx, err := ts.StartTransaction(ctx, sql.ReadWrite)
-			if err != nil {
-				return err
-			}
+	if ctx.GetTransaction() != nil {
+		return nil
+	}
 
-			ctx.SetTransaction(tx)
+	ctx.GetLogger().Tracef("beginning new transaction")
+	ts, ok := ctx.Session.(sql.TransactionSession)
+	if ok {
+		tx, err := ts.StartTransaction(ctx, sql.ReadWrite)
+		if err != nil {
+			return err
 		}
+
+		ctx.SetTransaction(tx)
 	}
 
 	return nil
@@ -782,6 +805,10 @@ func (e *Engine) EngineAnalyzer() *analyzer.Analyzer {
 	return e.Analyzer
 }
 
+func (e *Engine) EngineEventScheduler() sql.EventScheduler {
+	return e.EventScheduler
+}
+
 // InitializeEventScheduler initializes the EventScheduler for the engine with the given sql.Context
 // getter function, |ctxGetterFunc, the EventScheduler |status|, and the |period| for the event scheduler
 // to check for events to execute. If |period| is less than 1, then it is ignored and the default period
@@ -792,8 +819,6 @@ func (e *Engine) InitializeEventScheduler(ctxGetterFunc func() (*sql.Context, fu
 	if err != nil {
 		return err
 	}
-
-	e.Analyzer.EventScheduler = e.EventScheduler
 	return nil
 }
 
@@ -830,7 +855,7 @@ func (e *Engine) executeEvent(ctx *sql.Context, dbName, createEventStatement, us
 		return err
 	}
 
-	iter = finalizeIters(ctx, definitionNode, nil, iter)
+	iter, _ = rowexec.FinalizeIters(ctx, definitionNode, nil, iter)
 
 	// Drain the iterate to execute the event body/definition
 	// NOTE: No row data is returned for an event; we just need to execute the statements
@@ -863,12 +888,3 @@ func findCreateEventNode(planTree sql.Node) (*plan.CreateEvent, error) {
 
 	return createEventNode, nil
 }
-
-// finalizeIters applies the final transformations on sql.RowIter before execution.
-func finalizeIters(ctx *sql.Context, analyzed sql.Node, qFlags *sql.QueryFlags, iter sql.RowIter) sql.RowIter {
-	iter = rowexec.AddTriggerRollbackIter(ctx, qFlags, iter)
-	iter = rowexec.AddTransactionCommittingIter(qFlags, iter)
-	iter = plan.AddTrackedRowIter(ctx, analyzed, iter)
-	iter = rowexec.AddExpressionCloser(analyzed, iter)
-	return iter
-}

engine_test.go

@@ -231,7 +231,8 @@ func TestTrackProcess(t *testing.T) {
 	require.True(ok)
 
 	iter, err := rowexec.DefaultBuilder.Build(ctx, result, nil)
-	iter = finalizeIters(ctx, result, nil, iter)
+	require.NoError(err)
+	iter, _ = rowexec.FinalizeIters(ctx, result, nil, iter)
 	require.NoError(err)
 	_, err = sql.RowIterToRows(ctx, iter)
 	require.NoError(err)

enginetest/engine_only_test.go

@@ -442,7 +442,7 @@ func TestAnalyzer_Exp(t *testing.T) {
 			require.NoError(t, err)
 
 			ctx := enginetest.NewContext(harness)
-			b := planbuilder.New(ctx, e.EngineAnalyzer().Catalog, sql.NewMysqlParser())
+			b := planbuilder.New(ctx, e.EngineAnalyzer().Catalog, e.EngineEventScheduler(), nil)
 			parsed, _, _, _, err := b.Parse(tt.query, nil, false)
 			require.NoError(t, err)
 
@@ -866,10 +866,6 @@ func (s SimpleTableFunction) WithChildren(_ ...sql.Node) (sql.Node, error) {
 	return s, nil
 }
 
-func (s SimpleTableFunction) CheckPrivileges(_ *sql.Context, _ sql.PrivilegedOperationChecker) bool {
-	return true
-}
-
 // CollationCoercibility implements the interface sql.CollationCoercible.
 func (SimpleTableFunction) CollationCoercibility(ctx *sql.Context) (collation sql.CollationID, coercibility byte) {
 	return sql.Collation_binary, 7

enginetest/enginetests.go

@@ -583,7 +583,6 @@ func TestQueryPlan(t *testing.T, harness Harness, e QueryEngine, tt queries.Quer
 func TestQueryPlanWithName(t *testing.T, name string, harness Harness, e QueryEngine, query, expectedPlan string, options sql.DescribeOptions) {
 	t.Run(name, func(t *testing.T) {
 		ctx := NewContext(harness)
-
 		parsed, qFlags, err := planbuilder.Parse(ctx, e.EngineAnalyzer().Catalog, query)
 		require.NoError(t, err)
 
@@ -1710,12 +1709,12 @@ func TestInsertScriptsPrepared(t *testing.T, harness Harness) {
 func TestGeneratedColumns(t *testing.T, harness Harness) {
 	harness.Setup(setup.MydbData)
 	for _, script := range queries.GeneratedColumnTests {
-		TestScriptPrepared(t, harness, script)
+		TestScript(t, harness, script)
 	}
 	for _, script := range queries.BrokenGeneratedColumnTests {
 		t.Run(script.Name, func(t *testing.T) {
 			t.Skip(script.Name)
-			TestScriptPrepared(t, harness, script)
+			TestScript(t, harness, script)
 		})
 	}
 }

enginetest/evaluation.go

@@ -526,7 +526,7 @@ func injectBindVarsAndPrepare(
 		}
 	}
 
-	b := planbuilder.New(ctx, e.EngineAnalyzer().Catalog, sql.NewMysqlParser())
+	b := planbuilder.New(ctx, e.EngineAnalyzer().Catalog, e.EngineEventScheduler(), nil)
 	b.SetParserOptions(sql.LoadSqlMode(ctx).ParserOptions())
 	resPlan, _, err := b.BindOnly(parsed, q, nil)
 	if err != nil {

enginetest/memory_engine_test.go

@@ -1028,6 +1028,9 @@ func newMergableIndex(dbs []sql.Database, tableName string, exprs ...sql.Express
 	if db == nil {
 		return nil
 	}
+	if tableRevision, ok := table.(*memory.TableRevision); ok {
+		table = tableRevision.Table
+	}
 	return &memory.Index{
 		DB:         db.Name(),
 		DriverName: memory.IndexDriverId,

enginetest/plangen/cmd/plangen/main.go

@@ -165,7 +165,7 @@ func generatePlansForSuite(spec PlanSpec, w *bytes.Buffer) error {
 
 		if !tt.Skip {
 			ctx := enginetest.NewContext(harness)
-			binder := planbuilder.New(ctx, engine.EngineAnalyzer().Catalog, sql.NewMysqlParser())
+			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))