From 5dcdd2f4daf88a76cf571176f3d961b5997c08de Mon Sep 17 00:00:00 2001 From: Jille Timmermans Date: Tue, 27 Feb 2024 11:26:41 +0100 Subject: [PATCH 1/6] feat(compiler): Support subqueries in the FROM clause issue #2989, #2400 and probably others --- internal/compiler/query_catalog.go | 49 +++++++++++++++-- internal/compiler/resolve.go | 22 ++++++++ .../testdata/join_alias/mysql/go/query.sql.go | 54 +++++++++++++++++++ .../testdata/join_alias/mysql/query.sql | 6 +++ 4 files changed, 127 insertions(+), 4 deletions(-) diff --git a/internal/compiler/query_catalog.go b/internal/compiler/query_catalog.go index 80b59d876c..96b42be7f3 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 b1fbb1990e..9a82aa325c 100644 --- a/internal/compiler/resolve.go +++ b/internal/compiler/resolve.go @@ -80,6 +80,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 { 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 836f62c6d4..e6d565560b 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 9b087bcae7..82bb2e5cd5 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 <= ?; From 462e9b08a46a408e62d7524109da15a62bebd816 Mon Sep 17 00:00:00 2001 From: Jille Timmermans Date: Fri, 5 Apr 2024 00:10:50 +0200 Subject: [PATCH 2/6] fix a crash --- internal/compiler/resolve.go | 40 +++++++++++++++++++----------------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/internal/compiler/resolve.go b/internal/compiler/resolve.go index 9a82aa325c..49fe52e898 100644 --- a/internal/compiler/resolve.go +++ b/internal/compiler/resolve.go @@ -80,26 +80,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 qc != nil { + 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 + if err := indexTable(catalog.Table{ + Rel: f.Rel, + Columns: catCols, + }); err != nil { + return nil, err + } } } From b1772b5cd2bee779b3ae22943e9b8c4b6aec6030 Mon Sep 17 00:00:00 2001 From: Jille Timmermans Date: Fri, 5 Apr 2024 00:14:21 +0200 Subject: [PATCH 3/6] regen a test --- .../select_subquery_alias/postgresql/pgx/go/query.sql.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/internal/endtoend/testdata/select_subquery_alias/postgresql/pgx/go/query.sql.go b/internal/endtoend/testdata/select_subquery_alias/postgresql/pgx/go/query.sql.go index 8744642eb0..09ceb6e0b3 100644 --- a/internal/endtoend/testdata/select_subquery_alias/postgresql/pgx/go/query.sql.go +++ b/internal/endtoend/testdata/select_subquery_alias/postgresql/pgx/go/query.sql.go @@ -24,19 +24,19 @@ where amounts.last_balance < $2 ` type FindWalletsParams struct { - Column1 pgtype.Text - Column2 pgtype.Numeric + Type string + LastBalance int32 } type FindWalletsRow struct { ID int64 Address string Balance pgtype.Numeric - TotalBalance pgtype.Numeric + TotalBalance int64 } func (q *Queries) FindWallets(ctx context.Context, arg FindWalletsParams) ([]FindWalletsRow, error) { - rows, err := q.db.Query(ctx, findWallets, arg.Column1, arg.Column2) + rows, err := q.db.Query(ctx, findWallets, arg.Type, arg.LastBalance) if err != nil { return nil, err } From 18985896e04b36146a745c828151e3024b9ef1ca Mon Sep 17 00:00:00 2001 From: Simon Klee Date: Sun, 25 Feb 2024 00:55:53 +0100 Subject: [PATCH 4/6] 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 --- internal/compiler/resolve.go | 33 +++++++++++++- .../testdata/cte_resolve_ref/issue.md | 2 + .../cte_resolve_ref/postgresql/pgx/go/db.go | 32 ++++++++++++++ .../postgresql/pgx/go/models.go | 13 ++++++ .../postgresql/pgx/go/query.sql.go | 44 +++++++++++++++++++ .../cte_resolve_ref/postgresql/pgx/query.sql | 21 +++++++++ .../cte_resolve_ref/postgresql/pgx/schema.sql | 9 ++++ .../cte_resolve_ref/postgresql/pgx/sqlc.yaml | 10 +++++ 8 files changed, 163 insertions(+), 1 deletion(-) create mode 100644 internal/endtoend/testdata/cte_resolve_ref/issue.md create mode 100644 internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/db.go create mode 100644 internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/models.go create mode 100644 internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/go/query.sql.go create mode 100644 internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/query.sql create mode 100644 internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/schema.sql create mode 100644 internal/endtoend/testdata/cte_resolve_ref/postgresql/pgx/sqlc.yaml diff --git a/internal/compiler/resolve.go b/internal/compiler/resolve.go index b1fbb1990e..dba2e7cecc 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 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 0000000000..9268a7d1f8 --- /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 0000000000..7c111fd4e7 --- /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 0000000000..9099b926b3 --- /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 0000000000..d4fb8b9c5f --- /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 0000000000..39c87dfc31 --- /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 0000000000..3029ac89be --- /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 0000000000..01489e0ffc --- /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 From 7d1b0c648ee30d22d698109bb2e229303043494e Mon Sep 17 00:00:00 2001 From: Jille Timmermans Date: Fri, 5 Apr 2024 19:15:58 +0000 Subject: [PATCH 5/6] Manually update test to mismatch `sqlc generate` but match the CI's desire :( --- .../select_subquery_alias/postgresql/pgx/go/query.sql.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/endtoend/testdata/select_subquery_alias/postgresql/pgx/go/query.sql.go b/internal/endtoend/testdata/select_subquery_alias/postgresql/pgx/go/query.sql.go index 09ceb6e0b3..0cb154dba8 100644 --- a/internal/endtoend/testdata/select_subquery_alias/postgresql/pgx/go/query.sql.go +++ b/internal/endtoend/testdata/select_subquery_alias/postgresql/pgx/go/query.sql.go @@ -25,14 +25,14 @@ where amounts.last_balance < $2 type FindWalletsParams struct { Type string - LastBalance int32 + LastBalance pgtype.Numeric } type FindWalletsRow struct { ID int64 Address string Balance pgtype.Numeric - TotalBalance int64 + TotalBalance pgtype.Numeric } func (q *Queries) FindWallets(ctx context.Context, arg FindWalletsParams) ([]FindWalletsRow, error) { From ecddc20157a683e95e8fcbcf2a0a2419349e3c59 Mon Sep 17 00:00:00 2001 From: Jille Timmermans Date: Fri, 5 Apr 2024 19:22:14 +0000 Subject: [PATCH 6/6] Updated the improved tests I've checked each of them, and it seems correct those fields lose their nullability --- .../testdata/cte_left_join/postgresql/pgx/go/query.sql.go | 6 ++---- .../cte_recursive_union/postgresql/pgx/go/query.sql.go | 4 +--- .../testdata/cte_update/postgresql/pgx/go/query.sql.go | 6 +++--- 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/internal/endtoend/testdata/cte_left_join/postgresql/pgx/go/query.sql.go b/internal/endtoend/testdata/cte_left_join/postgresql/pgx/go/query.sql.go index 271ac35651..24f9ce475b 100644 --- a/internal/endtoend/testdata/cte_left_join/postgresql/pgx/go/query.sql.go +++ b/internal/endtoend/testdata/cte_left_join/postgresql/pgx/go/query.sql.go @@ -7,8 +7,6 @@ package querytest import ( "context" - - "github.com/jackc/pgx/v5/pgtype" ) const badQuery = `-- name: BadQuery :exec @@ -28,7 +26,7 @@ FROM WHERE c1.name = $1 ` -func (q *Queries) BadQuery(ctx context.Context, dollar_1 pgtype.Text) error { - _, err := q.db.Exec(ctx, badQuery, dollar_1) +func (q *Queries) BadQuery(ctx context.Context, name string) error { + _, err := q.db.Exec(ctx, badQuery, name) return err } diff --git a/internal/endtoend/testdata/cte_recursive_union/postgresql/pgx/go/query.sql.go b/internal/endtoend/testdata/cte_recursive_union/postgresql/pgx/go/query.sql.go index aa74cb5221..9b3b44647a 100644 --- a/internal/endtoend/testdata/cte_recursive_union/postgresql/pgx/go/query.sql.go +++ b/internal/endtoend/testdata/cte_recursive_union/postgresql/pgx/go/query.sql.go @@ -7,8 +7,6 @@ package querytest import ( "context" - - "github.com/jackc/pgx/v5/pgtype" ) const listCaseIntentHistory = `-- name: ListCaseIntentHistory :many @@ -36,7 +34,7 @@ type ListCaseIntentHistoryRow struct { Group string } -func (q *Queries) ListCaseIntentHistory(ctx context.Context, caseIntentID pgtype.Int8) ([]ListCaseIntentHistoryRow, error) { +func (q *Queries) ListCaseIntentHistory(ctx context.Context, caseIntentID int64) ([]ListCaseIntentHistoryRow, error) { rows, err := q.db.Query(ctx, listCaseIntentHistory, caseIntentID) if err != nil { return nil, err diff --git a/internal/endtoend/testdata/cte_update/postgresql/pgx/go/query.sql.go b/internal/endtoend/testdata/cte_update/postgresql/pgx/go/query.sql.go index 1176c852ae..fe4b0e3b94 100644 --- a/internal/endtoend/testdata/cte_update/postgresql/pgx/go/query.sql.go +++ b/internal/endtoend/testdata/cte_update/postgresql/pgx/go/query.sql.go @@ -23,9 +23,9 @@ from updated_attribute ` type UpdateAttributeParams struct { - FilterValue pgtype.Bool - Value pgtype.Text - ID pgtype.Int8 + FilterValue bool + Value string + ID int64 } type UpdateAttributeRow struct {