feat(compiler): Support subqueries in the FROM clause (#3227)

* 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

* feat(compiler): Support subqueries in the FROM clause

issue #2989, #2400 and probably others

---------

Co-authored-by: Simon Klee <[email protected]>

Author: Jille

internal/compiler/query_catalog.go

@@ -9,26 +9,36 @@ 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
 	case *ast.InsertStmt:
 		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
 }
 

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 {

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"

internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/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 <= ?;