Separate the scope of tables to allow removal of aliases

sqlc-dev · ovadbar · Aug 20, 2023 · Aug 20, 2023 · Aug 20, 2023 · Aug 20, 2023
commit 34c1a5498a8f34f86d9650416ec39560c7e3b7c0
diff --git a/internal/compiler/find_params.go b/internal/compiler/find_params.go
@@ -10,7 +10,7 @@ import (
 func findParameters(root ast.Node) ([]paramRef, error) {
 	refs := make([]paramRef, 0)
 	errors := make([]error, 0)
-	v := paramSearch{seen: make(map[int]struct{}), refs: &refs, errs: &errors}
+	v := paramSearch{seen: make(map[int]struct{}), refs: &refs, errs: &errors, rvs: &[]*ast.RangeVar{}}
 	astutils.Walk(v, root)
 	if len(*v.errs) > 0 {
 		problems := *v.errs
@@ -22,6 +22,7 @@ func findParameters(root ast.Node) ([]paramRef, error) {
 
 type paramRef struct {
 	parent ast.Node
+	rvs    []*ast.RangeVar
 	rv     *ast.RangeVar
 	ref    *ast.ParamRef
 	name   string // Named parameter support
@@ -31,6 +32,7 @@ type paramSearch struct {
 	parent   ast.Node
 	rangeVar *ast.RangeVar
 	refs     *[]paramRef
+	rvs      *[]*ast.RangeVar
 	seen     map[int]struct{}
 	errs     *[]error
 
@@ -58,6 +60,7 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 		return p
 	}
 
+	var reset bool
 	switch n := node.(type) {
 
 	case *ast.A_Expr:
@@ -70,6 +73,7 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 		p.parent = n.FuncCall
 
 	case *ast.DeleteStmt:
+		reset = true
 		if n.LimitCount != nil {
 			p.limitCount = n.LimitCount
 		}
@@ -78,7 +82,12 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 		p.parent = node
 
 	case *ast.InsertStmt:
+		reset = true
+		if n.Relation != nil {
+			*p.rvs = append(*p.rvs, n.Relation)
+		}
 		if s, ok := n.SelectStmt.(*ast.SelectStmt); ok {
+			*p.rvs = append(*p.rvs, toTables(s.FromClause)...)
 			for i, item := range s.TargetList.Items {
 				target, ok := item.(*ast.ResTarget)
 				if !ok {
@@ -92,7 +101,7 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 					*p.errs = append(*p.errs, fmt.Errorf("INSERT has more expressions than target columns"))
 					return p
 				}
-				*p.refs = append(*p.refs, paramRef{parent: n.Cols.Items[i], ref: ref, rv: n.Relation})
+				*p.refs = append(*p.refs, paramRef{parent: n.Cols.Items[i], ref: ref, rv: n.Relation, rvs: *p.rvs})
 				p.seen[ref.Location] = struct{}{}
 			}
 			for _, item := range s.ValuesLists.Items {
@@ -109,13 +118,16 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 						*p.errs = append(*p.errs, fmt.Errorf("INSERT has more expressions than target columns"))
 						return p
 					}
-					*p.refs = append(*p.refs, paramRef{parent: n.Cols.Items[i], ref: ref, rv: n.Relation})
+					*p.refs = append(*p.refs, paramRef{parent: n.Cols.Items[i], ref: ref, rv: n.Relation, rvs: *p.rvs})
 					p.seen[ref.Location] = struct{}{}
 				}
 			}
 		}
 
 	case *ast.UpdateStmt:
+		reset = true
+		*p.rvs = append(*p.rvs, toTables(n.FromClause)...)
+		*p.rvs = append(*p.rvs, toTables(n.Relations)...)
 		for _, item := range n.TargetList.Items {
 			target, ok := item.(*ast.ResTarget)
 			if !ok {
@@ -130,7 +142,7 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 				if !ok {
 					continue
 				}
-				*p.refs = append(*p.refs, paramRef{parent: target, ref: ref, rv: rv})
+				*p.refs = append(*p.refs, paramRef{parent: target, ref: ref, rv: rv, rvs: *p.rvs})
 			}
 			p.seen[ref.Location] = struct{}{}
 		}
@@ -139,12 +151,16 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 		}
 
 	case *ast.RangeVar:
+		if n != nil {
+			*p.rvs = append(*p.rvs, n)
+		}
 		p.rangeVar = n
 
 	case *ast.ResTarget:
 		p.parent = node
 
 	case *ast.SelectStmt:
+		reset = true
 		if n.LimitCount != nil {
 			p.limitCount = n.LimitCount
 		}
@@ -191,7 +207,7 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 		}
 
 		if set {
-			*p.refs = append(*p.refs, paramRef{parent: parent, ref: n, rv: p.rangeVar})
+			*p.refs = append(*p.refs, paramRef{parent: parent, ref: n, rv: p.rangeVar, rvs: *p.rvs})
 			p.seen[n.Location] = struct{}{}
 		}
 		return nil
@@ -215,5 +231,20 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 			p.Visit(n.Expr)
 		}
 	}
+	if reset {
+		rvs := *p.rvs
+		return paramSearch{seen: p.seen, refs: p.refs, errs: p.errs, rvs: &rvs, parent: p.parent, rangeVar: p.rangeVar, limitCount: p.limitCount, limitOffset: p.limitOffset}
+	}
 	return p
 }
+
+func toTables(tbl *ast.List) []*ast.RangeVar {
+	tables := make([]*ast.RangeVar, len(tbl.Items))
+	for _, t := range tbl.Items {
+		item, ok := t.(*ast.RangeVar)
+		if ok && item != nil {
+			tables = append(tables, item)
+		}
+	}
+	return tables
+}
diff --git a/internal/compiler/parse.go b/internal/compiler/parse.go
@@ -95,7 +95,11 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 		return nil, err
 	}
 
-	params, err := c.resolveCatalogRefs(qc, rvs, refs, namedParams, embeds)
+	err = c.resolveCatalogEmbeds(qc, rvs, embeds)
+	if err != nil {
+		return nil, err
+	}
+	params, err := c.resolveCatalogRefs(qc, refs, namedParams)
 	if err != nil {
 		return nil, err
 	}

diff --git a/internal/compiler/resolve.go b/internal/compiler/resolve.go
@@ -20,7 +20,7 @@ func dataType(n *ast.TypeName) string {
 	}
 }
 
-func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar, args []paramRef, params *named.ParamSet, embeds rewrite.EmbedSet) ([]Parameter, error) {
+func (comp *Compiler) resolveCatalogEmbeds(qc *QueryCatalog, rvs []*ast.RangeVar, embeds rewrite.EmbedSet) error {
 	c := comp.catalog
 
 	aliasMap := map[string]*ast.TableName{}
@@ -55,7 +55,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 		}
 		fqn, err := ParseTableName(rv)
 		if err != nil {
-			return nil, err
+			return err
 		}
 		if _, found := aliasMap[fqn.Name]; found {
 			continue
@@ -64,13 +64,13 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 		if err != nil {
 			// If the table name doesn't exist, fisrt check if it's a CTE
 			if _, qcerr := qc.GetTable(fqn); qcerr != nil {
-				return nil, err
+				return err
 			}
 			continue
 		}
 		err = indexTable(table)
 		if err != nil {
-			return nil, err
+			return err
 		}
 		if rv.Alias != nil {
 			aliasMap[*rv.Alias.Aliasname] = fqn
@@ -90,11 +90,71 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 			continue
 		}
 
-		return nil, fmt.Errorf("unable to resolve table with %q: %w", embed.Orig(), err)
+		return fmt.Errorf("unable to resolve table with %q: %w", embed.Orig(), err)
 	}
+	return nil
+}
 
+func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, args []paramRef, params *named.ParamSet) ([]Parameter, error) {
+	c := comp.catalog
+
+	// resolve a table for an embed
 	var a []Parameter
 	for _, ref := range args {
+		aliasMap := map[string]*ast.TableName{}
+		// TODO: Deprecate defaultTable
+		var defaultTable *ast.TableName
+		var tables []*ast.TableName
+
+		typeMap := map[string]map[string]map[string]*catalog.Column{}
+		indexTable := func(table catalog.Table) error {
+			tables = append(tables, table.Rel)
+			if defaultTable == nil {
+				defaultTable = table.Rel
+			}
+			schema := table.Rel.Schema
+			if schema == "" {
+				schema = c.DefaultSchema
+			}
+			if _, exists := typeMap[schema]; !exists {
+				typeMap[schema] = map[string]map[string]*catalog.Column{}
+			}
+			typeMap[schema][table.Rel.Name] = map[string]*catalog.Column{}
+			for _, c := range table.Columns {
+				cc := c
+				typeMap[schema][table.Rel.Name][c.Name] = cc
+			}
+			return nil
+		}
+
+		for _, rv := range ref.rvs {
+			if rv == nil || rv.Relname == nil {
+				continue
+			}
+			fqn, err := ParseTableName(rv)
+			if err != nil {
+				return nil, err
+			}
+			if _, found := aliasMap[fqn.Name]; found {
+				continue
+			}
+			table, err := c.GetTable(fqn)
+			if err != nil {
+				// If the table name doesn't exist, fisrt check if it's a CTE
+				if _, qcerr := qc.GetTable(fqn); qcerr != nil {
+					return nil, err
+				}
+				continue
+			}
+			err = indexTable(table)
+			if err != nil {
+				return nil, err
+			}
+			if rv.Alias != nil {
+				aliasMap[*rv.Alias.Aliasname] = fqn
+			}
+		}
+
 		switch n := ref.parent.(type) {
 
 		case *limitOffset:
@@ -196,7 +256,12 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 				}
 
 				var found int
+				seenTable := make(map[string]bool, len(search))
 				for _, table := range search {
+					if seenTable[table.Name] {
+						continue
+					}
+					seenTable[table.Name] = true
 					schema := table.Schema
 					if schema == "" {
 						schema = c.DefaultSchema
@@ -236,6 +301,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					}
 				}
 				if found > 1 {
+					fmt.Println("ambiguous 3")
 					return nil, &sqlerr.Error{
 						Code:     "42703",
 						Message:  fmt.Sprintf("column reference %q is ambiguous", key),
@@ -551,7 +617,12 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					}
 				}
 
+				seenTables := make(map[string]bool, len(search))
 				for _, table := range search {
+					if seenTables[table.Name] {
+						continue
+					}
+					seenTables[table.Name] = true
 					schema := table.Schema
 					if schema == "" {
 						schema = c.DefaultSchema

diff --git a/internal/endtoend/testdata/subquery_with_where/go/db.go b/internal/endtoend/testdata/subquery_with_where/go/db.go
diff --git a/internal/endtoend/testdata/subquery_with_where/go/models.go b/internal/endtoend/testdata/subquery_with_where/go/models.go
diff --git a/internal/endtoend/testdata/subquery_with_where/go/query.sql.go b/internal/endtoend/testdata/subquery_with_where/go/query.sql.go