merge with main

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/server.go

@@ -37,14 +37,19 @@ type ProtocolListener interface {
 }
 
 // ProtocolListenerFunc returns a ProtocolListener based on the configuration it was given.
-type ProtocolListenerFunc func(cfg mysql.ListenerConfig, sel ServerEventListener) (ProtocolListener, error)
-
-// DefaultProtocolListenerFunc is the protocol listener, which defaults to Vitess' protocol listener. Changing
-// this function will change the protocol listener used when creating all servers. If multiple servers are needed
-// with different protocols, then create each server after changing this function. Servers retain the protocol that
-// they were created with.
-var DefaultProtocolListenerFunc ProtocolListenerFunc = func(cfg mysql.ListenerConfig, sel ServerEventListener) (ProtocolListener, error) {
-	return mysql.NewListenerWithConfig(cfg)
+type ProtocolListenerFunc func(cfg Config, listenerCfg mysql.ListenerConfig, sel ServerEventListener) (ProtocolListener, error)
+
+func MySQLProtocolListenerFactory(cfg Config, listenerCfg mysql.ListenerConfig, sel ServerEventListener) (ProtocolListener, error) {
+	vtListener, err := mysql.NewListenerWithConfig(listenerCfg)
+	if err != nil {
+		return nil, err
+	}
+	if cfg.Version != "" {
+		vtListener.ServerVersion = cfg.Version
+	}
+	vtListener.TLSConfig = cfg.TLSConfig
+	vtListener.RequireSecureTransport = cfg.RequireSecureTransport
+	return vtListener, nil
 }
 
 type ServerEventListener interface {
@@ -114,10 +119,6 @@ func portInUse(hostPort string) bool {
 }
 
 func newServerFromHandler(cfg Config, e *sqle.Engine, sm *SessionManager, handler mysql.Handler, sel ServerEventListener) (*Server, error) {
-	for _, option := range cfg.Options {
-		option(e, sm, handler)
-	}
-
 	if cfg.ConnReadTimeout < 0 {
 		cfg.ConnReadTimeout = 0
 	}
@@ -128,6 +129,10 @@ func newServerFromHandler(cfg Config, e *sqle.Engine, sm *SessionManager, handle
 		cfg.MaxConnections = 0
 	}
 
+	for _, opt := range cfg.Options {
+		e, sm, handler = opt(e, sm, handler)
+	}
+
 	l := cfg.Listener
 	var unixSocketInUse error
 	if l == nil {
@@ -156,19 +161,15 @@ func newServerFromHandler(cfg Config, e *sqle.Engine, sm *SessionManager, handle
 		ConnReadBufferSize:       mysql.DefaultConnBufferSize,
 		AllowClearTextWithoutTLS: cfg.AllowClearTextWithoutTLS,
 	}
-	protocolListener, err := DefaultProtocolListenerFunc(listenerCfg, sel)
+	plf := cfg.ProtocolListenerFactory
+	if plf == nil {
+		plf = MySQLProtocolListenerFactory
+	}
+	protocolListener, err := plf(cfg, listenerCfg, sel)
 	if err != nil {
 		return nil, err
 	}
 
-	if vtListener, ok := protocolListener.(*mysql.Listener); ok {
-		if cfg.Version != "" {
-			vtListener.ServerVersion = cfg.Version
-		}
-		vtListener.TLSConfig = cfg.TLSConfig
-		vtListener.RequireSecureTransport = cfg.RequireSecureTransport
-	}
-
 	return &Server{
 		Listener:   protocolListener,
 		handler:    handler,

server/server_config.go

@@ -23,13 +23,9 @@ import (
 	"go.opentelemetry.io/otel/trace"
 
 	gms "github.com/dolthub/go-mysql-server"
-	sqle "github.com/dolthub/go-mysql-server"
 	"github.com/dolthub/go-mysql-server/sql"
 )
 
-// Option is an option to customize server.
-type Option func(e *sqle.Engine, sm *SessionManager, handler mysql.Handler)
-
 // Server is a MySQL server for SQLe engines.
 type Server struct {
 	Listener   ProtocolListener
@@ -38,6 +34,9 @@ type Server struct {
 	Engine     *gms.Engine
 }
 
+// An option to customize the server.
+type Option func(e *gms.Engine, sm *SessionManager, handler mysql.Handler) (*gms.Engine, *SessionManager, mysql.Handler)
+
 // Config for the mysql server.
 type Config struct {
 	// Protocol for the connection.
@@ -82,8 +81,14 @@ type Config struct {
 	// If true, queries will be logged as base64 encoded strings.
 	// If false (default behavior), queries will be logged as strings, but newlines and tabs will be replaced with spaces.
 	EncodeLoggedQuery bool
-	// Options add additional options to customize the server.
+	// Options gets a chance to visit and mutate the GMS *Engine,
+	// *server.SessionManager and the mysql.Handler as the server
+	// is being initialized, before the ProtocolListener is
+	// constructed.
 	Options []Option
+	// Used to get the ProtocolListener on server start.
+	// If unset, defaults to MySQLProtocolListenerFactory.
+	ProtocolListenerFactory ProtocolListenerFunc
 }
 
 func (c Config) NewConfig() (Config, error) {

sql/analyzer/analyzer.go

@@ -16,6 +16,7 @@ package analyzer
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"reflect"
 	"runtime/trace"
@@ -176,6 +177,10 @@ func (ab *Builder) RemoveAfterAllRule(id RuleId) *Builder {
 
 var log = logrus.New()
 
+func SetOutput(w io.Writer) {
+	log.SetOutput(w)
+}
+
 func init() {
 	// TODO: give the option for debug analyzer logging format to match the global one
 	log.SetFormatter(simpleLogFormatter{})

sql/analyzer/costed_index_scan.go

@@ -82,7 +82,7 @@ func costedIndexScans(ctx *sql.Context, a *Analyzer, n sql.Node, qFlags *sql.Que
 			}
 		}
 		if iat, ok := rt.UnderlyingTable().(sql.IndexAddressableTable); ok {
-			return costedIndexLookup(ctx, n, a.Catalog, iat, rt, aliasName, filter.Expression, qFlags)
+			return costedIndexLookup(ctx, n, a, iat, rt, aliasName, filter.Expression, qFlags)
 		}
 		return n, transform.SameTree, nil
 	})
@@ -123,19 +123,24 @@ func indexSearchableLookup(n sql.Node, rt sql.TableNode, lookup sql.IndexLookup,
 	return ret, transform.NewTree, nil
 }
 
-func costedIndexLookup(ctx *sql.Context, n sql.Node, cat sql.Catalog, iat sql.IndexAddressableTable, rt sql.TableNode, aliasName string, oldFilter sql.Expression, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
+func costedIndexLookup(ctx *sql.Context, n sql.Node, a *Analyzer, iat sql.IndexAddressableTable, rt sql.TableNode, aliasName string, oldFilter sql.Expression, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
 	indexes, err := iat.GetIndexes(ctx)
 	if err != nil {
 		return n, transform.SameTree, err
 	}
-	ita, _, filters, err := getCostedIndexScan(ctx, cat, rt, indexes, expression.SplitConjunction(oldFilter), qFlags)
+	ita, stats, filters, err := getCostedIndexScan(ctx, a.Catalog, rt, indexes, expression.SplitConjunction(oldFilter), qFlags)
 	if err != nil || ita == nil {
 		return n, transform.SameTree, err
 	}
 	var ret sql.Node = ita
 	if aliasName != "" {
 		ret = plan.NewTableAlias(aliasName, ret)
 	}
+
+	a.Log("new indexed table: %s/%s/%s", ita.Index().Database(), ita.Index().Table(), ita.Index().ID())
+	a.Log("index stats cnt: %d", stats.RowCount())
+	a.Log("index stats histogram: %s", stats.Histogram().DebugString())
+
 	// excluded from tree + not included in index scan => filter above scan
 	if len(filters) > 0 {
 		ret = plan.NewFilter(expression.JoinAnd(filters...), ret)