fix(resolve): fix resolving reference to CTEs

Fixed resolving refs to CTEs by adding CTEs to the aliasMap and indexing
its columns when resolving catalog references.

Fix #3219

Author: simonklee

internal/compiler/resolve.go

@@ -67,7 +67,38 @@ 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 {
+			catTable, qcerr := qc.GetTable(fqn)
+			if qcerr != nil {
+				return nil, err
+			}
+
+			// If it's a CTE, add it to the alias map and add its columns to
+			// the type map. This is to allow us to resolve references to the
+			// CTE's columns in a query.
+			aliasMap[fqn.Name] = fqn
+			if rv.Alias != nil {
+				aliasMap[*rv.Alias.Aliasname] = fqn
+			}
+
+			catCols := make([]*catalog.Column, 0, len(catTable.Columns))
+			for _, col := range catTable.Columns {
+				catCols = append(catCols, &catalog.Column{
+					Name:       col.Name,
+					Type:       ast.TypeName{Name: col.DataType},
+					IsNotNull:  col.NotNull,
+					IsUnsigned: col.Unsigned,
+					IsArray:    col.IsArray,
+					ArrayDims:  col.ArrayDims,
+					Comment:    col.Comment,
+					Length:     col.Length,
+				})
+			}
+
+			err = indexTable(catalog.Table{
+				Rel:     catTable.Rel,
+				Columns: catCols,
+			})
+			if err != nil {
 				return nil, err
 			}
 			continue

internal/endtoend/testdata/cte_resolve_ref/issue.md

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

internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/db.go

@@ -0,0 +1,21 @@
+-- name: CTERef :one
+WITH t1_ids AS (
+    SELECT id FROM t1
+)
+SELECT * FROM t1_ids WHERE t1_ids.id = sqlc.arg('id');
+
+-- name: CTEMultipleRefs :one
+WITH t1_ids AS (
+    SELECT id FROM t1 WHERE t1.id = sqlc.arg('id')
+),
+t2_ids AS (
+    SELECT id FROM t2 WHERE t2.id = sqlc.arg('id')
+),
+all_ids AS (
+    SELECT * FROM t1_ids
+    UNION
+    SELECT * FROM t2_ids
+)
+SELECT * FROM all_ids AS ids WHERE ids.id = sqlc.arg('id');
+
+

internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/models.go

@@ -0,0 +1,9 @@
+CREATE TABLE t1
+(
+    id SERIAL NOT NULL PRIMARY KEY
+);
+
+CREATE TABLE t2
+(
+    id SERIAL NOT NULL PRIMARY KEY
+);

internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/query.sql.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"