fix: resolver fails to look up CTE columns

Closes #4228.

The resolver fails to infer the type of a CTE column because it only
searches normal tables. My solution is to have two passes, one where
we search the CTE columns first (making them shadow regular tables) then
if not found, we search through regular tables.

We cannot simply add the CTE column in the existing `typeMap` because a
CTE and a regular table will have the same column name which will
trigger the `found > 1` check and the column will be considered
ambiguous. So we use a different typeMap for CTEs so we can separate the
search

Author: msaher

internal/compiler/resolve.go

@@ -28,7 +28,9 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 	// TODO: Deprecate defaultTable
 	var defaultTable *ast.TableName
 	var tables []*ast.TableName
+	var ctes []*ast.TableName
 
+	typeMapCte := map[string]map[string]map[string]*catalog.Column{}
 	typeMap := map[string]map[string]map[string]*catalog.Column{}
 	indexTable := func(table catalog.Table) error {
 		tables = append(tables, table.Rel)
@@ -67,9 +69,35 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 				continue
 			}
 			// If the table name doesn't exist, first check if it's a CTE
-			if _, qcerr := qc.GetTable(fqn); qcerr != nil {
+			cte, qcerr := qc.GetTable(fqn)
+			if qcerr != nil {
 				return nil, err
 			}
+			// duplicated logic from indexTable()
+			schema := fqn.Schema
+			if schema == "" {
+				schema = c.DefaultSchema
+			}
+			if _, exists := typeMapCte[schema]; !exists {
+				typeMapCte[schema] = map[string]map[string]*catalog.Column{}
+			}
+			typeMapCte[schema][fqn.Name] = map[string]*catalog.Column{}
+			for _, col := range cte.Columns {
+				cc := &catalog.Column{
+					Name:       col.Name,
+					IsNotNull:  col.NotNull,
+					IsUnsigned: col.Unsigned,
+					IsArray:    col.IsArray,
+					ArrayDims:  col.ArrayDims,
+					Comment:    col.Comment,
+					Length:     col.Length,
+				}
+				if col.Type != nil {
+					cc.Type = *col.Type
+				}
+				typeMapCte[schema][fqn.Name][col.Name] = cc
+			}
+			ctes = append(ctes, fqn)
 			continue
 		}
 		err = indexTable(table)
@@ -195,7 +223,61 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					panic("too many field items: " + strconv.Itoa(len(items)))
 				}
 
-				search := tables
+				search := ctes
+				if alias != "" {
+				    if _, ok := aliasMap[alias]; ok {
+				        // alias maps to a real table, skip CTE search entirely
+				        search = []*ast.TableName{}
+				    }
+				}
+
+				// resolve using ctes first
+				var found int
+				for _, table := range search {
+					schema := table.Schema
+					if schema == "" {
+						schema = c.DefaultSchema
+					}
+					if c, ok := typeMapCte[schema][table.Name][key]; ok {
+						found += 1
+						if ref.name != "" {
+							key = ref.name
+						}
+
+						defaultP := named.NewInferredParam(key, c.IsNotNull)
+						p, isNamed := params.FetchMerge(ref.ref.Number, defaultP)
+						a = append(a, Parameter{
+							Number: ref.ref.Number,
+							Column: &Column{
+								Name:         p.Name(),
+								OriginalName: c.Name,
+								DataType:     dataType(&c.Type),
+								NotNull:      p.NotNull(),
+								Unsigned:     c.IsUnsigned,
+								IsArray:      c.IsArray,
+								ArrayDims:    c.ArrayDims,
+								Length:       c.Length,
+								Table:        table,
+								IsNamedParam: isNamed,
+								IsSqlcSlice:  p.IsSqlcSlice(),
+							},
+						})
+					}
+				}
+
+				if found == 1 {
+					continue
+				}
+				if found > 1 {
+					return nil, &sqlerr.Error{
+						Code:     "42703",
+						Message:  fmt.Sprintf("column reference %q is ambiguous", key),
+						Location: node.Location,
+					}
+				}
+
+
+				search = tables
 				if alias != "" {
 					if original, ok := aliasMap[alias]; ok {
 						search = []*ast.TableName{original}
@@ -217,7 +299,9 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					}
 				}
 
-				var found int
+
+				// resolve using regular tables
+				found = 0
 				for _, table := range search {
 					schema := table.Schema
 					if schema == "" {
@@ -286,6 +370,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					schema = c.DefaultSchema
 				}
 
+				// should we use also look in CTEs?
 				if c, ok := typeMap[schema][table.Name][key]; ok {
 					defaultP := named.NewInferredParam(key, c.IsNotNull)
 					p, isNamed := params.FetchMerge(ref.ref.Number, defaultP)
@@ -475,6 +560,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 				schema = c.DefaultSchema
 			}
 
+			// should we use also look in CTEs?
 			tableMap, ok := typeMap[schema][rel]
 			if !ok {
 				return nil, sqlerr.RelationNotFound(rel)
@@ -583,6 +669,7 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					}
 				}
 
+				// should we use also look in CTEs?
 				for _, table := range search {
 					schema := table.Schema
 					if schema == "" {

internal/endtoend/testdata/cte_left_join/postgresql/pgx/go/query.sql.go

@@ -0,0 +1 @@
+https://github.com/sqlc-dev/sqlc/issues/4288

internal/endtoend/testdata/cte_renamed_column/issue.md

@@ -0,0 +1,5 @@
+-- name: GetUser :one
+WITH found AS (
+    SELECT id, id AS id2 FROM users WHERE id = $1
+)
+SELECT id2 + $2 AS result FROM found;

internal/endtoend/testdata/cte_renamed_column/postgresql/go/db.go

@@ -0,0 +1,5 @@
+-- Schema
+CREATE TABLE users (
+    id   INT PRIMARY KEY,
+    name TEXT NOT NULL
+);

internal/endtoend/testdata/cte_renamed_column/postgresql/go/models.go

@@ -0,0 +1,10 @@
+version: "2"
+sql:
+  - engine: "postgresql"
+    schema: "schema.sql"
+    queries: "query.sql"
+    gen:
+      go:
+        package: "querytest"
+        out: "go"
+        sql_package: "pgx/v5"