Merge branch 'main' into zachmu/enginetests5

Author: zachmu

engine.go

@@ -454,7 +454,7 @@ func (e *Engine) QueryWithBindings(ctx *sql.Context, query string, parsed sqlpar
 
 // PrepQueryPlanForExecution prepares a query plan for execution and returns the result schema with a row iterator to
 // begin spooling results
-func (e *Engine) PrepQueryPlanForExecution(ctx *sql.Context, _ string, plan sql.Node) (sql.Schema, sql.RowIter, *sql.QueryFlags, error) {
+func (e *Engine) PrepQueryPlanForExecution(ctx *sql.Context, _ string, plan sql.Node, qFlags *sql.QueryFlags) (sql.Schema, sql.RowIter, *sql.QueryFlags, error) {
 	// Give the integrator a chance to reject the session before proceeding
 	// TODO: this check doesn't belong here
 	err := ctx.Session.ValidateSession(ctx)
@@ -482,9 +482,9 @@ func (e *Engine) PrepQueryPlanForExecution(ctx *sql.Context, _ string, plan sql.
 		return nil, nil, nil, err
 	}
 
-	iter = finalizeIters(ctx, plan, nil, iter)
+	iter = finalizeIters(ctx, plan, qFlags, iter)
 
-	return plan.Schema(), iter, nil, nil
+	return plan.Schema(), iter, qFlags, nil
 }
 
 // BoundQueryPlan returns query plan for the given statement with the given bindings applied
@@ -866,6 +866,7 @@ func findCreateEventNode(planTree sql.Node) (*plan.CreateEvent, error) {
 
 // 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)

enginetest/queries/integration_plans.go

@@ -372,15 +372,15 @@ func (h *Handler) doQuery(
 	bindings map[string]*querypb.BindVariable,
 	callback func(*sqltypes.Result, bool) error,
 	qFlags *sql.QueryFlags,
-) (string, error) {
-	sqlCtx, err := h.sm.NewContext(ctx, c, query)
+) (remainder string, err error) {
+	var sqlCtx *sql.Context
+	sqlCtx, err = h.sm.NewContext(ctx, c, query)
 	if err != nil {
 		return "", err
 	}
 
 	start := time.Now()
 
-	var remainder string
 	var prequery string
 	if parsed == nil {
 		_, inPreparedCache := h.e.PreparedDataCache.GetCachedStmt(sqlCtx.Session.ID(), query)
@@ -411,23 +411,24 @@ func (h *Handler) doQuery(
 	sqlCtx.GetLogger().Debugf("Starting query")
 
 	finish := observeQuery(sqlCtx, query)
-	defer finish(err)
+	defer func() {
+		finish(err)
+	}()
 
 	sqlCtx.GetLogger().Tracef("beginning execution")
 
-	oCtx := ctx
-
 	// TODO: it would be nice to put this logic in the engine, not the handler, but we don't want the process to be
 	//  marked done until we're done spooling rows over the wire
-	ctx, err = sqlCtx.ProcessList.BeginQuery(sqlCtx, query)
-	defer func() {
-		if err != nil && ctx != nil {
-			sqlCtx.ProcessList.EndQuery(sqlCtx)
-		}
-	}()
+	sqlCtx, err = sqlCtx.ProcessList.BeginQuery(sqlCtx, query)
+	if err != nil {
+		return remainder, err
+	}
+	defer sqlCtx.ProcessList.EndQuery(sqlCtx)
 
+	var schema sql.Schema
+	var rowIter sql.RowIter
 	qFlags.Set(sql.QFlagDeferProjections)
-	schema, rowIter, qFlags, err := queryExec(sqlCtx, query, parsed, analyzedPlan, bindings, qFlags)
+	schema, rowIter, qFlags, err = queryExec(sqlCtx, query, parsed, analyzedPlan, bindings, qFlags)
 	if err != nil {
 		sqlCtx.GetLogger().WithError(err).Warn("error running query")
 		if verboseErrorLogging {
@@ -455,9 +456,6 @@ func (h *Handler) doQuery(
 		return remainder, err
 	}
 
-	// errGroup context is now canceled
-	ctx = oCtx
-
 	if err = setConnStatusFlags(sqlCtx, c); err != nil {
 		return remainder, err
 	}
@@ -1084,9 +1082,9 @@ func (h *Handler) executeBoundPlan(
 	_ sqlparser.Statement,
 	plan sql.Node,
 	_ map[string]*querypb.BindVariable,
-	_ *sql.QueryFlags,
+	qFlags *sql.QueryFlags,
 ) (sql.Schema, sql.RowIter, *sql.QueryFlags, error) {
-	return h.e.PrepQueryPlanForExecution(ctx, query, plan)
+	return h.e.PrepQueryPlanForExecution(ctx, query, plan, qFlags)
 }
 
 func bindingsToExprs(bindings map[string]*querypb.BindVariable) (map[string]sqlparser.Expr, error) {

server/handler.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/dolthub/vitess/go/mysql"
+	"github.com/dolthub/vitess/go/race"
 	"github.com/dolthub/vitess/go/sqltypes"
 	"github.com/dolthub/vitess/go/vt/proto/query"
 	"github.com/stretchr/testify/assert"
@@ -742,6 +743,113 @@ func TestHandlerKill(t *testing.T) {
 	require.Len(handler.sm.sessions, 1)
 }
 
+func TestHandlerKillQuery(t *testing.T) {
+	if race.Enabled {
+		t.Skip("this test is inherently racey")
+	}
+	require := require.New(t)
+	e, pro := setupMemDB(require)
+	dbFunc := pro.Database
+
+	handler := &Handler{
+		e: e,
+		sm: NewSessionManager(
+			func(ctx context.Context, conn *mysql.Conn, addr string) (sql.Session, error) {
+				return sql.NewBaseSessionWithClientServer(addr, sql.Client{Capabilities: conn.Capabilities}, conn.ConnectionID), nil
+			},
+			sql.NoopTracer,
+			dbFunc,
+			e.MemoryManager,
+			e.ProcessList,
+			"foo",
+		),
+	}
+
+	var err error
+	conn1 := newConn(1)
+	handler.NewConnection(conn1)
+
+	conn2 := newConn(2)
+	handler.NewConnection(conn2)
+
+	require.Len(handler.sm.connections, 2)
+	require.Len(handler.sm.sessions, 0)
+
+	handler.ComInitDB(conn1, "test")
+	err = handler.sm.SetDB(conn1, "test")
+	require.NoError(err)
+
+	err = handler.sm.SetDB(conn2, "test")
+	require.NoError(err)
+
+	require.False(conn1.Conn.(*mockConn).closed)
+	require.False(conn2.Conn.(*mockConn).closed)
+	require.Len(handler.sm.connections, 2)
+	require.Len(handler.sm.sessions, 2)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	sleepQuery := "SELECT SLEEP(1)"
+	go func() {
+		defer wg.Done()
+		err = handler.ComQuery(context.Background(), conn1, sleepQuery, func(res *sqltypes.Result, more bool) error {
+			return nil
+		})
+		require.Error(err)
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+	var sleepQueryID string
+	err = handler.ComQuery(context.Background(), conn2, "SHOW PROCESSLIST", func(res *sqltypes.Result, more bool) error {
+		// 1,  ,  , test, Query, 0, ...    , SELECT SLEEP(1000)
+		// 2,  ,  , test, Query, 0, running, SHOW PROCESSLIST
+		require.Equal(2, len(res.Rows))
+		hasSleepQuery := false
+		for _, row := range res.Rows {
+			if row[7].ToString() != sleepQuery {
+				continue
+			}
+			hasSleepQuery = true
+			sleepQueryID = row[0].ToString()
+			require.Equal("Query", row[4].ToString())
+		}
+		require.True(hasSleepQuery)
+		return nil
+	})
+	require.NoError(err)
+
+	time.Sleep(100 * time.Millisecond)
+	err = handler.ComQuery(context.Background(), conn2, "KILL QUERY "+sleepQueryID, func(res *sqltypes.Result, more bool) error {
+		return nil
+	})
+	require.NoError(err)
+	wg.Wait()
+
+	time.Sleep(100 * time.Millisecond)
+	err = handler.ComQuery(context.Background(), conn2, "SHOW PROCESSLIST", func(res *sqltypes.Result, more bool) error {
+		// 1,  ,  , test, Sleep, 0,        ,
+		// 2,  ,  , test, Query, 0, running, SHOW PROCESSLIST
+		require.Equal(2, len(res.Rows))
+		hasSleepQueryID := false
+		for _, row := range res.Rows {
+			if row[0].ToString() != sleepQueryID {
+				continue
+			}
+			hasSleepQueryID = true
+			require.Equal("Sleep", row[4].ToString())
+			require.Equal("", row[7].ToString())
+		}
+		require.True(hasSleepQueryID)
+		return nil
+	})
+	require.NoError(err)
+
+	require.False(conn1.Conn.(*mockConn).closed)
+	require.False(conn2.Conn.(*mockConn).closed)
+	require.Len(handler.sm.connections, 2)
+	require.Len(handler.sm.sessions, 2)
+}
+
 func TestSchemaToFields(t *testing.T) {
 	require := require.New(t)
 

server/handler_test.go

@@ -68,7 +68,6 @@ const (
 	assignRoutinesId             // assignRoutines
 	modifyUpdateExprsForJoinId   // modifyUpdateExprsForJoin
 	applyUpdateAccumulatorsId    // applyUpdateAccumulators
-	wrapWithRollbackId           // wrapWithRollback
 	applyForeignKeysId           // applyForeignKeys
 
 	// validate

sql/analyzer/rule_ids.go

@@ -22,7 +22,6 @@ func init() {
 		{applyTriggersId, applyTriggers},
 		{applyProceduresId, applyProcedures},
 		{applyUpdateAccumulatorsId, applyUpdateAccumulators},
-		{wrapWithRollbackId, wrapWithRollback},
 		{inlineSubqueryAliasRefsId, inlineSubqueryAliasRefs},
 		{cacheSubqueryAliasesInJoinsId, cacheSubqueryAliasesInJoins},
 		{BacktickDefaulColumnValueNamesId, backtickDefaultColumnValueNames},

sql/analyzer/ruleid_string.go

@@ -346,6 +346,7 @@ func applyTrigger(ctx *sql.Context, a *Analyzer, originalNode, n sql.Node, scope
 
 		switch n := c.Node.(type) {
 		case *plan.InsertInto:
+			qFlags.Set(sql.QFlagTrigger)
 			if trigger.TriggerTime == sqlparser.BeforeStr {
 				triggerExecutor := plan.NewTriggerExecutor(n.Source, triggerLogic, plan.InsertTrigger, plan.TriggerTime(trigger.TriggerTime), sql.TriggerDefinition{
 					Name:            trigger.TriggerName,
@@ -359,6 +360,7 @@ func applyTrigger(ctx *sql.Context, a *Analyzer, originalNode, n sql.Node, scope
 				}), transform.NewTree, nil
 			}
 		case *plan.Update:
+			qFlags.Set(sql.QFlagTrigger)
 			if trigger.TriggerTime == sqlparser.BeforeStr {
 				triggerExecutor := plan.NewTriggerExecutor(n.Child, triggerLogic, plan.UpdateTrigger, plan.TriggerTime(trigger.TriggerTime), sql.TriggerDefinition{
 					Name:            trigger.TriggerName,
@@ -387,6 +389,7 @@ func applyTrigger(ctx *sql.Context, a *Analyzer, originalNode, n sql.Node, scope
 					"does not support triggers; retry with single table deletes")
 			}
 
+			qFlags.Set(sql.QFlagTrigger)
 			if trigger.TriggerTime == sqlparser.BeforeStr {
 				triggerExecutor := plan.NewTriggerExecutor(n.Child, triggerLogic, plan.DeleteTrigger, plan.TriggerTime(trigger.TriggerTime), sql.TriggerDefinition{
 					Name:            trigger.TriggerName,
@@ -517,41 +520,3 @@ func orderTriggersAndReverseAfter(triggers []*plan.CreateTrigger) []*plan.Create
 func triggerEventsMatch(event plan.TriggerEvent, event2 string) bool {
 	return strings.ToLower((string)(event)) == strings.ToLower(event2)
 }
-
-// wrapWithRollback wraps the entire tree iff it contains a trigger, allowing rollback when a trigger errors
-func wrapWithRollback(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope, sel RuleSelector, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
-	// Check if tree contains a TriggerExecutor
-	containsTrigger := false
-	transform.Inspect(n, func(n sql.Node) bool {
-		// After Triggers wrap nodes
-		if _, ok := n.(*plan.TriggerExecutor); ok {
-			containsTrigger = true
-			return false // done, don't bother to recurse
-		}
-
-		// Before Triggers on Inserts are inside Source
-		if n, ok := n.(*plan.InsertInto); ok {
-			if _, ok := n.Source.(*plan.TriggerExecutor); ok {
-				containsTrigger = true
-				return false
-			}
-		}
-
-		// Before Triggers on Delete and Update should be in children
-		return true
-	})
-
-	// No TriggerExecutor, so return same tree
-	if !containsTrigger {
-		return n, transform.SameTree, nil
-	}
-
-	// If we don't have a transaction session we can't do rollbacks
-	_, ok := ctx.Session.(sql.TransactionSession)
-	if !ok {
-		return plan.NewNoopTriggerRollback(n), transform.NewTree, nil
-	}
-
-	// Wrap tree with new node
-	return plan.NewTriggerRollback(n), transform.NewTree, nil
-}