Skip to content

Commit 3f129e7

Browse files
Jillesimonklee
Jille
and
simonklee
authored
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]>
1 parent d227d5f commit 3f129e7

11 files changed

Lines changed: 290 additions & 5 deletions

File tree

‎internal/compiler/query_catalog.go‎

Lines changed: 45 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,26 +9,36 @@ import (
99
)
1010

1111
type QueryCatalog struct {
12-
catalog *catalog.Catalog
13-
ctes map[string]*Table
14-
embeds rewrite.EmbedSet
12+
catalog *catalog.Catalog
13+
ctes map[string]*Table
14+
fromClauses map[string]*Table
15+
embeds rewrite.EmbedSet
1516
}
1617

1718
func (comp *Compiler) buildQueryCatalog(c *catalog.Catalog, node ast.Node, embeds rewrite.EmbedSet) (*QueryCatalog, error) {
1819
var with *ast.WithClause
20+
var from *ast.List
1921
switch n := node.(type) {
2022
case *ast.DeleteStmt:
2123
with = n.WithClause
2224
case *ast.InsertStmt:
2325
with = n.WithClause
2426
case *ast.UpdateStmt:
2527
with = n.WithClause
28+
from = n.FromClause
2629
case *ast.SelectStmt:
2730
with = n.WithClause
31+
from = n.FromClause
2832
default:
2933
with = nil
34+
from = nil
35+
}
36+
qc := &QueryCatalog{
37+
catalog: c,
38+
ctes: map[string]*Table{},
39+
fromClauses: map[string]*Table{},
40+
embeds: embeds,
3041
}
31-
qc := &QueryCatalog{catalog: c, ctes: map[string]*Table{}, embeds: embeds}
3242
if with != nil {
3343
for _, item := range with.Ctes.Items {
3444
if cte, ok := item.(*ast.CommonTableExpr); ok {
@@ -60,6 +70,37 @@ func (comp *Compiler) buildQueryCatalog(c *catalog.Catalog, node ast.Node, embed
6070
}
6171
}
6272
}
73+
if from != nil {
74+
for _, item := range from.Items {
75+
if rs, ok := item.(*ast.RangeSubselect); ok {
76+
cols, err := comp.outputColumns(qc, rs.Subquery)
77+
if err != nil {
78+
return nil, err
79+
}
80+
var names []string
81+
if rs.Alias.Colnames != nil {
82+
for _, item := range rs.Alias.Colnames.Items {
83+
if val, ok := item.(*ast.String); ok {
84+
names = append(names, val.Str)
85+
} else {
86+
names = append(names, "")
87+
}
88+
}
89+
}
90+
rel := &ast.TableName{Name: *rs.Alias.Aliasname}
91+
for i := range cols {
92+
cols[i].Table = rel
93+
if len(names) > i {
94+
cols[i].Name = names[i]
95+
}
96+
}
97+
qc.fromClauses[*rs.Alias.Aliasname] = &Table{
98+
Rel: rel,
99+
Columns: cols,
100+
}
101+
}
102+
}
103+
}
63104
return qc, nil
64105
}
65106

‎internal/compiler/resolve.go‎

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,38 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
6767
continue
6868
}
6969
// If the table name doesn't exist, first check if it's a CTE
70-
if _, qcerr := qc.GetTable(fqn); qcerr != nil {
70+
catTable, qcerr := qc.GetTable(fqn)
71+
if qcerr != nil {
72+
return nil, err
73+
}
74+
75+
// If it's a CTE, add it to the alias map and add its columns to
76+
// the type map. This is to allow us to resolve references to the
77+
// CTE's columns in a query.
78+
aliasMap[fqn.Name] = fqn
79+
if rv.Alias != nil {
80+
aliasMap[*rv.Alias.Aliasname] = fqn
81+
}
82+
83+
catCols := make([]*catalog.Column, 0, len(catTable.Columns))
84+
for _, col := range catTable.Columns {
85+
catCols = append(catCols, &catalog.Column{
86+
Name: col.Name,
87+
Type: ast.TypeName{Name: col.DataType},
88+
IsNotNull: col.NotNull,
89+
IsUnsigned: col.Unsigned,
90+
IsArray: col.IsArray,
91+
ArrayDims: col.ArrayDims,
92+
Comment: col.Comment,
93+
Length: col.Length,
94+
})
95+
}
96+
97+
err = indexTable(catalog.Table{
98+
Rel: catTable.Rel,
99+
Columns: catCols,
100+
})
101+
if err != nil {
71102
return nil, err
72103
}
73104
continue
@@ -80,6 +111,28 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
80111
aliasMap[*rv.Alias.Aliasname] = fqn
81112
}
82113
}
114+
for _, f := range qc.fromClauses {
115+
catCols := make([]*catalog.Column, 0, len(f.Columns))
116+
for _, col := range f.Columns {
117+
catCols = append(catCols, &catalog.Column{
118+
Name: col.Name,
119+
Type: ast.TypeName{Name: col.DataType},
120+
IsNotNull: col.NotNull,
121+
IsUnsigned: col.Unsigned,
122+
IsArray: col.IsArray,
123+
ArrayDims: col.ArrayDims,
124+
Comment: col.Comment,
125+
Length: col.Length,
126+
})
127+
}
128+
129+
if err := indexTable(catalog.Table{
130+
Rel: f.Rel,
131+
Columns: catCols,
132+
}); err != nil {
133+
return nil, err
134+
}
135+
}
83136

84137
// resolve a table for an embed
85138
for _, embed := range embeds {

‎internal/endtoend/testdata/cte_resolve_ref/issue.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
https://github.com/sqlc-dev/sqlc/issues/3219
2+

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

Lines changed: 32 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 13 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/query.sql.go‎

Lines changed: 44 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/query.sql‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
-- name: CTERef :one
2+
WITH t1_ids AS (
3+
SELECT id FROM t1
4+
)
5+
SELECT * FROM t1_ids WHERE t1_ids.id = sqlc.arg('id');
6+
7+
-- name: CTEMultipleRefs :one
8+
WITH t1_ids AS (
9+
SELECT id FROM t1 WHERE t1.id = sqlc.arg('id')
10+
),
11+
t2_ids AS (
12+
SELECT id FROM t2 WHERE t2.id = sqlc.arg('id')
13+
),
14+
all_ids AS (
15+
SELECT * FROM t1_ids
16+
UNION
17+
SELECT * FROM t2_ids
18+
)
19+
SELECT * FROM all_ids AS ids WHERE ids.id = sqlc.arg('id');
20+
21+

‎internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/schema.sql‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
CREATE TABLE t1
2+
(
3+
id SERIAL NOT NULL PRIMARY KEY
4+
);
5+
6+
CREATE TABLE t2
7+
(
8+
id SERIAL NOT NULL PRIMARY KEY
9+
);

‎internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/sqlc.yaml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
version: "2"
2+
sql:
3+
- engine: "postgresql"
4+
schema: "schema.sql"
5+
queries: "query.sql"
6+
gen:
7+
go:
8+
package: "querytest"
9+
out: "go"
10+
sql_package: "pgx/v5"

‎internal/endtoend/testdata/join_alias/mysql/go/query.sql.go‎

Lines changed: 54 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)