wip - split analysis into two parts

Signed-off-by: Andres Taylor <andres@planetscale.com>

Author: systay

go/vt/vtgate/semantics/analyzer.go

@@ -28,43 +28,49 @@ import (
 // analyzer controls the flow of the analysis.
 // It starts the tree walking and controls which part of the analysis sees which parts of the tree
 type analyzer struct {
-	scoper   *scoper
-	tables   *tableCollector
-	binder   *binder
-	typer    *typer
-	rewriter *earlyRewriter
-	sig      QuerySignature
+	scoper      *scoper
+	earlyTables *earlyTableCollector
+	tables      *tableCollector
+	binder      *binder
+	typer       *typer
+	rewriter    *earlyRewriter
+	sig         QuerySignature
+	si          SchemaInformation
+	currentDb   string
 
 	err          error
 	inProjection int
 
-	projErr      error
-	unshardedErr error
-	warning      string
+	projErr                 error
+	unshardedErr            error
+	warning                 string
+	singleUnshardedKeyspace bool
 }
 
 // newAnalyzer create the semantic analyzer
 func newAnalyzer(dbName string, si SchemaInformation) *analyzer {
 	// TODO  dependencies between these components are a little tangled. We should try to clean up
 	s := newScoper()
 	a := &analyzer{
-		scoper: s,
-		tables: newTableCollector(s, si, dbName),
-		typer:  newTyper(si.Environment().CollationEnv()),
+		scoper:      s,
+		earlyTables: newEarlyTableCollector(si, dbName),
+		typer:       newTyper(si.Environment().CollationEnv()),
+		si:          si,
+		currentDb:   dbName,
 	}
 	s.org = a
-	a.tables.org = a
+	return a
+}
 
-	b := newBinder(s, a, a.tables, a.typer)
-	a.binder = b
+func (a *analyzer) lateInit() {
+	a.tables = a.earlyTables.newTableCollector(a.scoper, a)
+	a.binder = newBinder(a.scoper, a, a.tables, a.typer)
 	a.rewriter = &earlyRewriter{
-		env:             si.Environment(),
-		scoper:          s,
-		binder:          b,
+		env:             a.si.Environment(),
+		scoper:          a.scoper,
+		binder:          a.binder,
 		expandedColumns: map[sqlparser.TableName][]*sqlparser.ColName{},
 	}
-	s.binder = b
-	return a
 }
 
 // Analyze analyzes the parsed query.
@@ -109,6 +115,32 @@ func (a *analyzer) newSemTable(
 	if isCommented {
 		comments = commentedStmt.GetParsedComments()
 	}
+
+	if a.singleUnshardedKeyspace {
+		return &SemTable{
+			Tables:                    a.earlyTables.Tables,
+			Comments:                  comments,
+			Warning:                   a.warning,
+			Collation:                 coll,
+			ExprTypes:                 map[sqlparser.Expr]evalengine.Type{},
+			NotSingleRouteErr:         a.projErr,
+			NotUnshardedErr:           a.unshardedErr,
+			Recursive:                 ExprDependencies{},
+			Direct:                    ExprDependencies{},
+			Targets:                   map[sqlparser.IdentifierCS]TableSet{},
+			ColumnEqualities:          map[columnName][]sqlparser.Expr{},
+			ExpandedColumns:           map[sqlparser.TableName][]*sqlparser.ColName{},
+			columns:                   map[*sqlparser.Union]sqlparser.SelectExprs{},
+			comparator:                nil,
+			StatementIDs:              a.scoper.statementIDs,
+			QuerySignature:            QuerySignature{},
+			childForeignKeysInvolved:  map[TableSet][]vindexes.ChildFKInfo{},
+			parentForeignKeysInvolved: map[TableSet][]vindexes.ParentFKInfo{},
+			childFkToUpdExprs:         map[string]sqlparser.UpdateExprs{},
+			collEnv:                   env,
+		}, nil
+	}
+
 	columns := map[*sqlparser.Union]sqlparser.SelectExprs{}
 	for union, info := range a.tables.unionInfo {
 		columns[union] = info.exprs
@@ -298,10 +330,38 @@ func (a *analyzer) depsForExpr(expr sqlparser.Expr) (direct, recursive TableSet,
 }
 
 func (a *analyzer) analyze(statement sqlparser.Statement) error {
+	_ = sqlparser.Rewrite(statement, nil, a.earlyUp)
+	if a.err != nil {
+		return a.err
+	}
+	ks, _ := singleUnshardedKeyspace(a.earlyTables.Tables)
+	if ks != nil {
+		// if we found a single unsharded keyspace in the early walk we can stop here
+		a.singleUnshardedKeyspace = true
+		return nil
+	}
+
+	a.lateInit()
+
 	_ = sqlparser.Rewrite(statement, a.analyzeDown, a.analyzeUp)
 	return a.err
 }
 
+// earlyUp collects tables in the query, so we can check
+// if this a single unsharded query we are dealing with
+func (a *analyzer) earlyUp(cursor *sqlparser.Cursor) bool {
+	if !a.shouldContinue() {
+		return false
+	}
+
+	if err := a.earlyTables.up(cursor); err != nil {
+		a.setError(err)
+		return false
+	}
+
+	return a.shouldContinue()
+}
+
 func (a *analyzer) shouldContinue() bool {
 	return a.err == nil
 }

go/vt/vtgate/semantics/analyzer_test.go

@@ -986,7 +986,7 @@ func TestScopingWDerivedTables(t *testing.T) {
 			recursiveExpectation: MergeTableSets(TS0, TS1),
 		}, {
 			query:                "select t.id from (select * from user) as t join user as u on t.id = u.id",
-			expectation:          TS1,
+			expectation:          TS2,
 			recursiveExpectation: TS0,
 		}, {
 			query:                "select t.col1 from t3 ua join (select t1.id, t1.col1 from t1 join t2) as t",
@@ -1638,11 +1638,11 @@ func TestScopingSubQueryJoinClause(t *testing.T) {
 
 var ks1 = &vindexes.Keyspace{
 	Name:    "ks1",
-	Sharded: false,
+	Sharded: true,
 }
 var ks2 = &vindexes.Keyspace{
 	Name:    "ks2",
-	Sharded: false,
+	Sharded: true,
 }
 var ks3 = &vindexes.Keyspace{
 	Name:    "ks3",

go/vt/vtgate/semantics/semantic_state.go

@@ -730,6 +730,10 @@ func (st *SemTable) ColumnLookup(col *sqlparser.ColName) (int, error) {
 
 // SingleUnshardedKeyspace returns the single keyspace if all tables in the query are in the same, unsharded keyspace
 func (st *SemTable) SingleUnshardedKeyspace() (ks *vindexes.Keyspace, tables []*vindexes.Table) {
+	return singleUnshardedKeyspace(st.Tables)
+}
+
+func singleUnshardedKeyspace(tableInfos []TableInfo) (ks *vindexes.Keyspace, tables []*vindexes.Table) {
 	validKS := func(this *vindexes.Keyspace) bool {
 		if this == nil || this.Sharded {
 			return false
@@ -744,7 +748,7 @@ func (st *SemTable) SingleUnshardedKeyspace() (ks *vindexes.Keyspace, tables []*
 		return true
 	}
 
-	for _, table := range st.Tables {
+	for _, table := range tableInfos {
 		if _, isDT := table.(*DerivedTable); isDT {
 			continue
 		}

go/vt/vtgate/semantics/table_collector.go

@@ -35,15 +35,61 @@ type tableCollector struct {
 	currentDb string
 	org       originable
 	unionInfo map[*sqlparser.Union]unionInfo
+	done      map[*sqlparser.AliasedTableExpr]TableInfo
 }
 
-func newTableCollector(scoper *scoper, si SchemaInformation, currentDb string) *tableCollector {
-	return &tableCollector{
-		scoper:    scoper,
+type earlyTableCollector struct {
+	si        SchemaInformation
+	currentDb string
+	Tables    []TableInfo
+	done      map[*sqlparser.AliasedTableExpr]TableInfo
+}
+
+func newEarlyTableCollector(si SchemaInformation, currentDb string) *earlyTableCollector {
+	return &earlyTableCollector{
 		si:        si,
 		currentDb: currentDb,
+		done:      map[*sqlparser.AliasedTableExpr]TableInfo{},
+	}
+}
+
+func (etc *earlyTableCollector) up(cursor *sqlparser.Cursor) error {
+	aet, ok := cursor.Node().(*sqlparser.AliasedTableExpr)
+	if !ok {
+		return nil
+	}
+	return etc.visitAliasedTableExpr(aet)
+}
+
+func (etc *earlyTableCollector) visitAliasedTableExpr(aet *sqlparser.AliasedTableExpr) error {
+	tbl, ok := aet.Expr.(sqlparser.TableName)
+	if !ok {
+		return nil
+	}
+	return etc.handleTableName(tbl, aet)
+}
+
+func (etc *earlyTableCollector) newTableCollector(scoper *scoper, org originable) *tableCollector {
+	return &tableCollector{
+		Tables:    etc.Tables,
+		scoper:    scoper,
+		si:        etc.si,
+		currentDb: etc.currentDb,
 		unionInfo: map[*sqlparser.Union]unionInfo{},
+		done:      etc.done,
+		org:       org,
+	}
+}
+
+func (etc *earlyTableCollector) handleTableName(tbl sqlparser.TableName, aet *sqlparser.AliasedTableExpr) error {
+	tableInfo, err := getTableInfo(aet, tbl, etc.si, etc.currentDb)
+	if err != nil {
+		return err
 	}
+
+	etc.done[aet] = tableInfo
+	etc.Tables = append(etc.Tables, tableInfo)
+	return nil
 }
 
 func (tc *tableCollector) up(cursor *sqlparser.Cursor) error {
@@ -103,28 +149,42 @@ func (tc *tableCollector) visitAliasedTableExpr(node *sqlparser.AliasedTableExpr
 	return nil
 }
 
-func (tc *tableCollector) handleTableName(node *sqlparser.AliasedTableExpr, t sqlparser.TableName) error {
+func (tc *tableCollector) handleTableName(node *sqlparser.AliasedTableExpr, t sqlparser.TableName) (err error) {
+	var tableInfo TableInfo
+	var found bool
+
+	tableInfo, found = tc.done[node]
+	if !found {
+		tableInfo, err = getTableInfo(node, t, tc.si, tc.currentDb)
+		if err != nil {
+			return err
+		}
+		tc.Tables = append(tc.Tables, tableInfo)
+	}
+
+	scope := tc.scoper.currentScope()
+	return scope.addTable(tableInfo)
+}
+
+func getTableInfo(node *sqlparser.AliasedTableExpr, t sqlparser.TableName, si SchemaInformation, currentDb string) (TableInfo, error) {
 	var tbl *vindexes.Table
 	var vindex vindexes.Vindex
 	isInfSchema := sqlparser.SystemSchema(t.Qualifier.String())
 	var err error
-	tbl, vindex, _, _, _, err = tc.si.FindTableOrVindex(t)
+	tbl, vindex, _, _, _, err = si.FindTableOrVindex(t)
 	if err != nil && !isInfSchema {
 		// if we are dealing with a system table, it might not be available in the vschema, but that is OK
-		return err
+		return nil, err
 	}
 	if tbl == nil && vindex != nil {
 		tbl = newVindexTable(t.Name)
 	}
 
-	scope := tc.scoper.currentScope()
-	tableInfo, err := tc.createTable(t, node, tbl, isInfSchema, vindex)
+	tableInfo, err := createTable(t, node, tbl, isInfSchema, vindex, si, currentDb)
 	if err != nil {
-		return err
+		return nil, err
 	}
-
-	tc.Tables = append(tc.Tables, tableInfo)
-	return scope.addTable(tableInfo)
+	return tableInfo, nil
 }
 
 func (tc *tableCollector) handleDerivedTable(node *sqlparser.AliasedTableExpr, t *sqlparser.DerivedTable) error {
@@ -228,12 +288,14 @@ func (tc *tableCollector) tableInfoFor(id TableSet) (TableInfo, error) {
 	return tc.Tables[offset], nil
 }
 
-func (tc *tableCollector) createTable(
+func createTable(
 	t sqlparser.TableName,
 	alias *sqlparser.AliasedTableExpr,
 	tbl *vindexes.Table,
 	isInfSchema bool,
 	vindex vindexes.Vindex,
+	si SchemaInformation,
+	currentDb string,
 ) (TableInfo, error) {
 	hint := getVindexHint(alias.Hints)
 
@@ -247,13 +309,13 @@ func (tc *tableCollector) createTable(
 		Table:        tbl,
 		VindexHint:   hint,
 		isInfSchema:  isInfSchema,
-		collationEnv: tc.si.Environment().CollationEnv(),
+		collationEnv: si.Environment().CollationEnv(),
 	}
 
 	if alias.As.IsEmpty() {
 		dbName := t.Qualifier.String()
 		if dbName == "" {
-			dbName = tc.currentDb
+			dbName = currentDb
 		}
 
 		table.dbName = dbName