Passing context through

Author: fulghum

server/handler_test.go

@@ -38,7 +38,6 @@ import (
 )
 
 func TestHandlerOutput(t *testing.T) {
-
 	e := setupMemDB(require.New(t))
 	dummyConn := newConn(1)
 	handler := &Handler{
@@ -655,7 +654,26 @@ func TestSchemaToFields(t *testing.T) {
 
 	require.Equal(len(schema), len(expected))
 
-	fields := schemaToFields(schema)
+	handler := &Handler{
+		e: setupMemDB(require),
+		sm: NewSessionManager(
+			testSessionBuilder,
+			sql.NoopTracer,
+			func(ctx *sql.Context, db string) bool { return db == "test" },
+			sql.NewMemoryManager(nil),
+			sqle.NewProcessList(),
+			"foo",
+		),
+		readTimeout: time.Second,
+	}
+
+	conn := newConn(1)
+	handler.NewConnection(conn)
+
+	ctx, err := handler.sm.NewContextWithQuery(conn, "SELECT 1")
+	require.NoError(err)
+
+	fields := schemaToFields(ctx, schema)
 	for i := 0; i < len(fields); i++ {
 		t.Run(schema[i].Name, func(t *testing.T) {
 			assert.Equal(t, expected[i], fields[i])

sql/analyzer/validate_create_table.go

@@ -159,7 +159,7 @@ func resolveAlterColumn(ctx *sql.Context, a *Analyzer, n sql.Node, scope *Scope,
 			if err != nil {
 				return nil, transform.SameTree, err
 			}
-			sch, err = validatePrimaryKey(initialSch, sch, n.(*plan.AlterPK))
+			sch, err = validatePrimaryKey(ctx, initialSch, sch, n.(*plan.AlterPK))
 			if err != nil {
 				return nil, transform.SameTree, err
 			}
@@ -610,7 +610,7 @@ func validateIndexes(ctx *sql.Context, tableSpec *plan.TableSpec) error {
 			if !ok {
 				return sql.ErrUnknownIndexColumn.New(idxCol.Name, idx.IndexName)
 			}
-			err := validatePrefixLength(schCol, idxCol)
+			err := validatePrefixLength(ctx, schCol, idxCol)
 			if err != nil {
 				return err
 			}
@@ -683,7 +683,7 @@ func getTableIndexNames(ctx *sql.Context, a *Analyzer, table sql.Node) ([]string
 }
 
 // validatePrimaryKey validates a primary key add or drop operation.
-func validatePrimaryKey(initialSch, sch sql.Schema, ai *plan.AlterPK) (sql.Schema, error) {
+func validatePrimaryKey(ctx *sql.Context, initialSch, sch sql.Schema, ai *plan.AlterPK) (sql.Schema, error) {
 	tableName := getTableName(ai.Table)
 	switch ai.Action {
 	case plan.PrimaryKeyAction_Create:
@@ -698,7 +698,7 @@ func validatePrimaryKey(initialSch, sch sql.Schema, ai *plan.AlterPK) (sql.Schem
 
 		for _, idxCol := range ai.Columns {
 			schCol := sch[sch.IndexOf(idxCol.Name, tableName)]
-			err := validatePrefixLength(schCol, idxCol)
+			err := validatePrefixLength(ctx, schCol, idxCol)
 			if err != nil {
 				return nil, err
 			}