merge

Author: James Cor

enginetest/queries/information_schema_queries.go

@@ -874,6 +874,18 @@ var InfoSchemaScripts = []ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "issue 8930: connect to info schema",
+		SetUpScript: []string{
+			"use information_schema",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "SELECT C.COLUMN_NAME AS label, 'connection.column' as \"type\", C.TABLE_NAME AS \"table\", C.DATA_TYPE AS \"dataType\", CAST(C.CHARACTER_MAXIMUM_LENGTH AS UNSIGNED) AS size, CAST(UPPER( CONCAT( C.DATA_TYPE, CASE WHEN C.DATA_TYPE = 'text' THEN '' ELSE ( CASE WHEN C.CHARACTER_MAXIMUM_LENGTH > 0 THEN ( CONCAT('(', C.CHARACTER_MAXIMUM_LENGTH, ')') ) ELSE '' END ) END ) ) AS CHAR CHARACTER SET utf8) AS \"detail\", C.TABLE_CATALOG AS \"catalog\", C.TABLE_SCHEMA AS \"database\", C.TABLE_SCHEMA AS \"schema\", C.COLUMN_DEFAULT AS \"defaultValue\", C.IS_NULLABLE AS \"isNullable\", (CASE WHEN C.COLUMN_KEY = 'PRI' THEN 1 ELSE 0 END) AS \"isPk\", (CASE WHEN KCU.REFERENCED_COLUMN_NAME IS NULL THEN 0 ELSE 1 END) AS \"isFk\" FROM INFORMATION_SCHEMA.COLUMNS AS C LEFT JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS KCU ON ( C.TABLE_NAME = KCU.TABLE_NAME AND C.TABLE_SCHEMA = KCU.TABLE_SCHEMA AND C.TABLE_CATALOG = KCU.TABLE_CATALOG AND C.COLUMN_NAME = KCU.COLUMN_NAME ) JOIN INFORMATION_SCHEMA.TABLES AS T ON C.TABLE_NAME = T.TABLE_NAME AND C.TABLE_SCHEMA = T.TABLE_SCHEMA AND C.TABLE_CATALOG = T.TABLE_CATALOG WHERE C.TABLE_SCHEMA = 'dev' AND C.TABLE_NAME = 'countries' AND C.TABLE_CATALOG = 'def' ORDER BY C.TABLE_NAME, C.ORDINAL_POSITION",
+				Expected: []sql.Row{},
+			},
+		},
+	},
 	{
 		Name: "query does not use optimization rule on LIKE clause because info_schema db charset is utf8mb3",
 		SetUpScript: []string{

enginetest/server_engine_test.go

@@ -83,7 +83,7 @@ type serverScriptTestAssertion struct {
 	expectedRows         []any
 
 	// can't avoid writing custom comparator because of how gosql.Rows.Scan() works
-	checkRows func(rows *gosql.Rows, expectedRows []any) (bool, error)
+	checkRows func(t *testing.T, rows *gosql.Rows, expectedRows []any) (bool, error)
 }
 
 type serverScriptTest struct {
@@ -108,7 +108,7 @@ func TestServerPreparedStatements(t *testing.T) {
 					expectedRows: []any{
 						[]float64{321.4},
 					},
-					checkRows: func(rows *gosql.Rows, expectedRows []any) (bool, error) {
+					checkRows: func(t *testing.T, rows *gosql.Rows, expectedRows []any) (bool, error) {
 						var i float64
 						var rowNum int
 						for rows.Next() {
@@ -133,7 +133,7 @@ func TestServerPreparedStatements(t *testing.T) {
 						[]float64{213.4},
 						[]float64{213.4},
 					},
-					checkRows: func(rows *gosql.Rows, expectedRows []any) (bool, error) {
+					checkRows: func(t *testing.T, rows *gosql.Rows, expectedRows []any) (bool, error) {
 						var i float64
 						var rowNum int
 						for rows.Next() {
@@ -197,7 +197,7 @@ func TestServerPreparedStatements(t *testing.T) {
 						[]uint64{uint64(math.MaxInt64 + 1)},
 						[]uint64{uint64(math.MaxUint64)},
 					},
-					checkRows: func(rows *gosql.Rows, expectedRows []any) (bool, error) {
+					checkRows: func(t *testing.T, rows *gosql.Rows, expectedRows []any) (bool, error) {
 						var i uint64
 						var rowNum int
 						for rows.Next() {
@@ -247,7 +247,7 @@ func TestServerPreparedStatements(t *testing.T) {
 						[]int64{int64(-1)},
 						[]int64{int64(math.MaxInt64)},
 					},
-					checkRows: func(rows *gosql.Rows, expectedRows []any) (bool, error) {
+					checkRows: func(t *testing.T, rows *gosql.Rows, expectedRows []any) (bool, error) {
 						var i int64
 						var rowNum int
 						for rows.Next() {
@@ -275,12 +275,12 @@ func TestServerPreparedStatements(t *testing.T) {
 			},
 			assertions: []serverScriptTestAssertion{
 				{
-					query: "select * from test where c0 = 2 and c1 = 3;",
+					query: "select * from test where c0 = 2 and c1 = 3 order by pk;",
 					expectedRows: []any{
 						[]uint64{uint64(2), uint64(3), uint64(1)},
 						[]uint64{uint64(2), uint64(3), uint64(7)},
 					},
-					checkRows: func(rows *gosql.Rows, expectedRows []any) (bool, error) {
+					checkRows: func(t *testing.T, rows *gosql.Rows, expectedRows []any) (bool, error) {
 						var c0, c1, pk uint64
 						var rowNum int
 						for rows.Next() {
@@ -363,7 +363,7 @@ func TestServerPreparedStatements(t *testing.T) {
 					} else {
 						require.NoError(t, err)
 					}
-					ok, err := assertion.checkRows(rows, assertion.expectedRows)
+					ok, err := assertion.checkRows(t, rows, assertion.expectedRows)
 					require.NoError(t, err)
 					require.True(t, ok)
 				})

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/dolthub/go-icu-regex v0.0.0-20250303123116-549b8d7cad00
 	github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71
 	github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81
-	github.com/dolthub/vitess v0.0.0-20250303224041-5cc89c183bc4
+	github.com/dolthub/vitess v0.0.0-20250304211657-920ca9ec2b9a
 	github.com/go-kit/kit v0.10.0
 	github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d
 	github.com/gocraft/dbr/v2 v2.7.2

go.sum

@@ -62,6 +62,8 @@ github.com/dolthub/vitess v0.0.0-20250228011932-c4f6bba87730 h1:GtlMVB7+Z7fZZj7B
 github.com/dolthub/vitess v0.0.0-20250228011932-c4f6bba87730/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70=
 github.com/dolthub/vitess v0.0.0-20250303224041-5cc89c183bc4 h1:wtS9ZWEyEeYzLCcqdGUo+7i3hAV5MWuY9Z7tYbQa65A=
 github.com/dolthub/vitess v0.0.0-20250303224041-5cc89c183bc4/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70=
+github.com/dolthub/vitess v0.0.0-20250304211657-920ca9ec2b9a h1:HIH9g4z+yXr4DIFyT6L5qOIEGJ1zVtlj6baPyHAG4Yw=
+github.com/dolthub/vitess v0.0.0-20250304211657-920ca9ec2b9a/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70=
 github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/eapache/go-resiliency v1.1.0/go.mod h1:kFI+JgMyC7bLPUVY133qvEBtVayf5mFgVsvEsIPBvNs=
 github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21/go.mod h1:+020luEh2TKB4/GOp8oxxtq0Daoen/Cii55CzbTV6DU=

server/context.go

@@ -49,6 +49,10 @@ type SessionManager struct {
 	connections map[uint32]*mysql.Conn
 	lastPid     uint64
 	ctxFactory  sql.ContextFactory
+	// Implements WaitForClosedConnections(), which is only used
+	// at server shutdown to allow the integrator to ensure that
+	// no connections are being handled by handlers.
+	wg sync.WaitGroup
 }
 
 // NewSessionManager creates a SessionManager with the given SessionBuilder.
@@ -82,6 +86,13 @@ func (s *SessionManager) nextPid() uint64 {
 	return s.lastPid
 }
 
+// Block the calling thread until all known connections are closed. It
+// is an error to call this concurrently while the server might still
+// be accepting new connections.
+func (s *SessionManager) WaitForClosedConnections() {
+	s.wg.Wait()
+}
+
 // AddConn adds a connection to be tracked by the SessionManager. Should be called as
 // soon as possible after the server has accepted the connection. Results in
 // the connection being tracked by ProcessList and being available through
@@ -93,6 +104,7 @@ func (s *SessionManager) AddConn(conn *mysql.Conn) {
 	defer s.mu.Unlock()
 	s.connections[conn.ConnectionID] = conn
 	s.processlist.AddConnection(conn.ConnectionID, conn.RemoteAddr().String())
+	s.wg.Add(1)
 }
 
 // NewSession creates a Session for the given connection and saves it to the session pool.
@@ -270,6 +282,7 @@ func (s *SessionManager) KillConnection(connID uint32) error {
 func (s *SessionManager) RemoveConn(conn *mysql.Conn) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	s.wg.Done()
 	if cur, ok := s.sessions[conn.ConnectionID]; ok {
 		sql.SessionEnd(cur)
 	}

server/extension.go

@@ -27,24 +27,36 @@ import (
 	sqle "github.com/dolthub/go-mysql-server"
 )
 
-func Intercept(h Interceptor) {
-	inters = append(inters, h)
-	sort.Slice(inters, func(i, j int) bool { return inters[i].Priority() < inters[j].Priority() })
-}
-
-func WithChain() Option {
-	return func(e *sqle.Engine, sm *SessionManager, handler mysql.Handler) {
-		f := DefaultProtocolListenerFunc
-		DefaultProtocolListenerFunc = func(cfg mysql.ListenerConfig, sel ServerEventListener) (ProtocolListener, error) {
-			cfg.Handler = buildChain(cfg.Handler)
-			return f(cfg, sel)
-		}
-	}
+// InterceptorChain allows an integrator to build a chain of
+// |Interceptor| instances which will wrap and intercept the server's
+// mysql.Handler.
+//
+// Example usage:
+//
+// var ic InterceptorChain
+// ic.WithInterceptor(metricsInterceptor)
+// ic.WithInterceptor(authInterceptor)
+// server, err := NewServer(Config{ ..., Options: []Option{ic.Option()}, ...}, ...)
+type InterceptorChain struct {
+	inters []Interceptor
 }
 
-var inters []Interceptor
+func (ic *InterceptorChain) WithInterceptor(h Interceptor) {
+	ic.inters = append(ic.inters, h)
+}
 
-func buildChain(h mysql.Handler) mysql.Handler {
+func (ic *InterceptorChain) Option() Option {
+	return func(e *sqle.Engine, sm *SessionManager, handler mysql.Handler) (*sqle.Engine, *SessionManager, mysql.Handler) {
+		chainHandler := buildChain(handler, ic.inters)
+		return e, sm, chainHandler
+	}
+}
+
+func buildChain(h mysql.Handler, inters []Interceptor) mysql.Handler {
+	// XXX: Mutates |inters|
+	sort.Slice(inters, func(i, j int) bool {
+		return inters[i].Priority() < inters[j].Priority()
+	})
 	var last Chain = h
 	for i := len(inters) - 1; i >= 0; i-- {
 		filter := inters[i]
@@ -55,7 +67,6 @@ func buildChain(h mysql.Handler) mysql.Handler {
 }
 
 type Interceptor interface {
-
 	// Priority returns the priority of the interceptor.
 	Priority() int
 
@@ -88,7 +99,6 @@ type Interceptor interface {
 }
 
 type Chain interface {
-
 	// ComQuery is called when a connection receives a query.
 	// Note the contents of the query slice may change after
 	// the first call to callback. So the Handler should not

server/handler.go

@@ -17,6 +17,7 @@ package server
 import (
 	"context"
 	"encoding/base64"
+	goerrors "errors"
 	"fmt"
 	"io"
 	"net"
@@ -609,31 +610,32 @@ func resultForMax1RowIter(ctx *sql.Context, schema sql.Schema, iter sql.RowIter,
 
 // resultForDefaultIter reads batches of rows from the iterator
 // and writes results into the callback function.
-func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema sql.Schema, iter sql.RowIter, callback func(*sqltypes.Result, bool) error, resultFields []*querypb.Field, more bool, buf *sql.ByteBuffer) (r *sqltypes.Result, processedAtLeastOneBatch bool, returnErr error) {
+func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema sql.Schema, iter sql.RowIter, callback func(*sqltypes.Result, bool) error, resultFields []*querypb.Field, more bool, buf *sql.ByteBuffer) (*sqltypes.Result, bool, error) {
 	defer trace.StartRegion(ctx, "Handler.resultForDefaultIter").End()
 
 	eg, ctx := ctx.NewErrgroup()
-
-	pan2err := func() {
+	pan2err := func(err *error) {
 		if recoveredPanic := recover(); recoveredPanic != nil {
-			returnErr = fmt.Errorf("handler caught panic: %v", recoveredPanic)
+			*err = goerrors.Join(*err, fmt.Errorf("handler caught panic: %v", recoveredPanic))
 		}
 	}
-
 	wg := sync.WaitGroup{}
 	wg.Add(2)
 
+	var r *sqltypes.Result
+	var processedAtLeastOneBatch bool
+
 	// Read rows off the row iterator and send them to the row channel.
 	iter, projs := GetDeferredProjections(iter)
 	var rowChan = make(chan sql.Row, 512)
-	eg.Go(func() error {
-		defer pan2err()
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
 		defer wg.Done()
 		defer close(rowChan)
 		for {
 			select {
 			case <-ctx.Done():
-				return nil
+				return context.Cause(ctx)
 			default:
 				row, err := iter.Next(ctx)
 				if err == io.EOF {
@@ -651,9 +653,12 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 		}
 	})
 
+	// TODO: poll for closed connections should obviously also run even if
+	// we're doing something with an OK result or a single row result, etc.
+	// This should be in the caller.
 	pollCtx, cancelF := ctx.NewSubContext()
-	eg.Go(func() error {
-		defer pan2err()
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
 		return h.pollForClosedConnection(pollCtx, c)
 	})
 
@@ -676,8 +681,8 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 
 	// Reads rows from the channel, converts them to wire format,
 	// and calls |callback| to give them to vitess.
-	eg.Go(func() error {
-		defer pan2err()
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
 		defer cancelF()
 		defer wg.Done()
 		for {
@@ -695,7 +700,7 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 
 			select {
 			case <-ctx.Done():
-				return nil
+				return context.Cause(ctx)
 			case row, ok := <-rowChan:
 				if !ok {
 					return nil
@@ -716,6 +721,9 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 				ctx.GetLogger().Tracef("spooling result row %s", outputRow)
 				r.Rows = append(r.Rows, outputRow)
 				r.RowsAffected++
+				if !timer.Stop() {
+					<-timer.C
+				}
 			case <-timer.C:
 				// TODO: timer should probably go in its own thread, as rowChan is blocking
 				if h.readTimeout != 0 {
@@ -724,17 +732,14 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 					return ErrRowTimeout.New()
 				}
 			}
-			if !timer.Stop() {
-				<-timer.C
-			}
 			timer.Reset(waitTime)
 		}
 	})
 
 	// Close() kills this PID in the process list,
 	// wait until all rows have be sent over the wire
-	eg.Go(func() error {
-		defer pan2err()
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
 		wg.Wait()
 		return iter.Close(ctx)
 	})
@@ -745,9 +750,9 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 		if verboseErrorLogging {
 			fmt.Printf("Err: %+v", err)
 		}
-		returnErr = err
+		return nil, false, err
 	}
-	return
+	return r, processedAtLeastOneBatch, nil
 }
 
 // See https://dev.mysql.com/doc/internals/en/status-flags.html