Repository: sqlc-dev/sqlc
Pull request #3310 (2024-04-03 06:57:42)

feat(compiler): Support subqueries in the FROM clause (second coming)
issue https://github.com/sqlc-dev/sqlc/issues/2989, https://github.com/sqlc-dev/sqlc/issues/2400 and probably others

This is the same as https://github.com/sqlc-dev/sqlc/pull/3227, but rebased to main on Kyle's request.

simonklee: Thanks for rebasing, @Jille. 

Patch:
diff --git a/internal/compiler/query_catalog.go b/internal/compiler/query_catalog.go
index 80b59d876..96b42be7f 100644
--- a/internal/compiler/query_catalog.go
+++ b/internal/compiler/query_catalog.go
@@ -9,13 +9,15 @@ import (
 )
 
 type QueryCatalog struct {
-	catalog *catalog.Catalog
-	ctes    map[string]*Table
-	embeds  rewrite.EmbedSet
+	catalog     *catalog.Catalog
+	ctes        map[string]*Table
+	fromClauses map[string]*Table
+	embeds      rewrite.EmbedSet
 }
 
 func (comp *Compiler) buildQueryCatalog(c *catalog.Catalog, node ast.Node, embeds rewrite.EmbedSet) (*QueryCatalog, error) {
 	var with *ast.WithClause
+	var from *ast.List
 	switch n := node.(type) {
 	case *ast.DeleteStmt:
 		with = n.WithClause
@@ -23,12 +25,20 @@ func (comp *Compiler) buildQueryCatalog(c *catalog.Catalog, node ast.Node, embed
 		with = n.WithClause
 	case *ast.UpdateStmt:
 		with = n.WithClause
+		from = n.FromClause
 	case *ast.SelectStmt:
 		with = n.WithClause
+		from = n.FromClause
 	default:
 		with = nil
+		from = nil
+	}
+	qc := &QueryCatalog{
+		catalog:     c,
+		ctes:        map[string]*Table{},
+		fromClauses: map[string]*Table{},
+		embeds:      embeds,
 	}
-	qc := &QueryCatalog{catalog: c, ctes: map[string]*Table{}, embeds: embeds}
 	if with != nil {
 		for _, item := range with.Ctes.Items {
 			if cte, ok := item.(*ast.CommonTableExpr); ok {
@@ -60,6 +70,37 @@ func (comp *Compiler) buildQueryCatalog(c *catalog.Catalog, node ast.Node, embed
 			}
 		}
 	}
+	if from != nil {
+		for _, item := range from.Items {
+			if rs, ok := item.(*ast.RangeSubselect); ok {
+				cols, err := comp.outputColumns(qc, rs.Subquery)
+				if err != nil {
+					return nil, err
+				}
+				var names []string
+				if rs.Alias.Colnames != nil {
+					for _, item := range rs.Alias.Colnames.Items {
+						if val, ok := item.(*ast.String); ok {
+							names = append(names, val.Str)
+						} else {
+							names = append(names, "")
+						}
+					}
+				}
+				rel := &ast.TableName{Name: *rs.Alias.Aliasname}
+				for i := range cols {
+					cols[i].Table = rel
+					if len(names) > i {
+						cols[i].Name = names[i]
+					}
+				}
+				qc.fromClauses[*rs.Alias.Aliasname] = &Table{
+					Rel:     rel,
+					Columns: cols,
+				}
+			}
+		}
+	}
 	return qc, nil
 }
 
diff --git a/internal/compiler/resolve.go b/internal/compiler/resolve.go
index b1fbb1990..a001a9ace 100644
--- a/internal/compiler/resolve.go
+++ b/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
@@ -80,6 +111,28 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 			aliasMap[*rv.Alias.Aliasname] = fqn
 		}
 	}
+	for _, f := range qc.fromClauses {
+		catCols := make([]*catalog.Column, 0, len(f.Columns))
+		for _, col := range f.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,
+			})
+		}
+
+		if err := indexTable(catalog.Table{
+			Rel:     f.Rel,
+			Columns: catCols,
+		}); err != nil {
+			return nil, err
+		}
+	}
 
 	// resolve a table for an embed
 	for _, embed := range embeds {


Test patch:
diff --git a/internal/endtoend/testdata/cte_resolve_ref/issue.md b/internal/endtoend/testdata/cte_resolve_ref/issue.md
new file mode 100644
index 000000000..9268a7d1f
--- /dev/null
+++ b/internal/endtoend/testdata/cte_resolve_ref/issue.md
@@ -0,0 +1,2 @@
+https://github.com/sqlc-dev/sqlc/issues/3219
+
diff --git a/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/db.go b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/db.go
new file mode 100644
index 000000000..7c111fd4e
--- /dev/null
+++ b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/db.go
@@ -0,0 +1,32 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.26.0
+
+package querytest
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+)
+
+type DBTX interface {
+	Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error)
+	Query(context.Context, string, ...interface{}) (pgx.Rows, error)
+	QueryRow(context.Context, string, ...interface{}) pgx.Row
+}
+
+func New(db DBTX) *Queries {
+	return &Queries{db: db}
+}
+
+type Queries struct {
+	db DBTX
+}
+
+func (q *Queries) WithTx(tx pgx.Tx) *Queries {
+	return &Queries{
+		db: tx,
+	}
+}
diff --git a/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/models.go b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/models.go
new file mode 100644
index 000000000..9099b926b
--- /dev/null
+++ b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/models.go
@@ -0,0 +1,13 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.26.0
+
+package querytest
+
+type T1 struct {
+	ID int32
+}
+
+type T2 struct {
+	ID int32
+}
diff --git a/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/query.sql.go b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/query.sql.go
new file mode 100644
index 000000000..d4fb8b9c5
--- /dev/null
+++ b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/query.sql.go
@@ -0,0 +1,44 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.26.0
+// source: query.sql
+
+package querytest
+
+import (
+	"context"
+)
+
+const cTEMultipleRefs = `-- name: CTEMultipleRefs :one
+WITH t1_ids AS (
+    SELECT id FROM t1 WHERE t1.id = $1
+),
+t2_ids AS (
+    SELECT id FROM t2 WHERE t2.id = $1
+),
+all_ids AS (
+    SELECT id FROM t1_ids
+    UNION
+    SELECT id FROM t2_ids
+)
+SELECT id FROM all_ids AS ids WHERE ids.id = $1
+`
+
+func (q *Queries) CTEMultipleRefs(ctx context.Context, id int32) (int32, error) {
+	row := q.db.QueryRow(ctx, cTEMultipleRefs, id)
+	err := row.Scan(&id)
+	return id, err
+}
+
+const cTERef = `-- name: CTERef :one
+WITH t1_ids AS (
+    SELECT id FROM t1
+)
+SELECT id FROM t1_ids WHERE t1_ids.id = $1
+`
+
+func (q *Queries) CTERef(ctx context.Context, id int32) (int32, error) {
+	row := q.db.QueryRow(ctx, cTERef, id)
+	err := row.Scan(&id)
+	return id, err
+}
diff --git a/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/query.sql b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/query.sql
new file mode 100644
index 000000000..39c87dfc3
--- /dev/null
+++ b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/query.sql
@@ -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');
+
+
diff --git a/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/schema.sql b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/schema.sql
new file mode 100644
index 000000000..3029ac89b
--- /dev/null
+++ b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/schema.sql
@@ -0,0 +1,9 @@
+CREATE TABLE t1
+(
+    id SERIAL NOT NULL PRIMARY KEY
+);
+
+CREATE TABLE t2
+(
+    id SERIAL NOT NULL PRIMARY KEY
+);
diff --git a/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/sqlc.yaml b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/sqlc.yaml
new file mode 100644
index 000000000..01489e0ff
--- /dev/null
+++ b/internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/sqlc.yaml
@@ -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"
\ No newline at end of file
diff --git a/internal/endtoend/testdata/join_alias/mysql/go/query.sql.go b/internal/endtoend/testdata/join_alias/mysql/go/query.sql.go
index 836f62c6d..e6d565560 100644
--- a/internal/endtoend/testdata/join_alias/mysql/go/query.sql.go
+++ b/internal/endtoend/testdata/join_alias/mysql/go/query.sql.go
@@ -80,3 +80,57 @@ func (q *Queries) AliasJoin(ctx context.Context, id uint64) ([]AliasJoinRow, err
 	}
 	return items, nil
 }
+
+const columnAlias = `-- name: ColumnAlias :many
+SELECT n FROM (SELECT 1 AS n) AS x WHERE n <= ?
+`
+
+func (q *Queries) ColumnAlias(ctx context.Context, n int32) ([]int32, error) {
+	rows, err := q.db.QueryContext(ctx, columnAlias, n)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []int32
+	for rows.Next() {
+		var n int32
+		if err := rows.Scan(&n); err != nil {
+			return nil, err
+		}
+		items = append(items, n)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
+const subqueryAlias = `-- name: SubqueryAlias :many
+SELECT n FROM (SELECT 1 AS n) AS x WHERE x.n <= ?
+`
+
+func (q *Queries) SubqueryAlias(ctx context.Context, n int32) ([]int32, error) {
+	rows, err := q.db.QueryContext(ctx, subqueryAlias, n)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []int32
+	for rows.Next() {
+		var n int32
+		if err := rows.Scan(&n); err != nil {
+			return nil, err
+		}
+		items = append(items, n)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
diff --git a/internal/endtoend/testdata/join_alias/mysql/query.sql b/internal/endtoend/testdata/join_alias/mysql/query.sql
index 9b087bcae..82bb2e5cd 100644
--- a/internal/endtoend/testdata/join_alias/mysql/query.sql
+++ b/internal/endtoend/testdata/join_alias/mysql/query.sql
@@ -9,3 +9,9 @@ SELECT *
 FROM foo f
 JOIN bar b ON b.id = f.id
 WHERE f.id = ?;
+
+-- name: SubqueryAlias :many
+SELECT * FROM (SELECT 1 AS n) AS x WHERE x.n <= ?;
+
+-- name: ColumnAlias :many
+SELECT * FROM (SELECT 1 AS n) AS x WHERE n <= ?;