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..4b7184a242 --- /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.25.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..b1e92d036a --- /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.25.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..c6c4a019bb --- /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.25.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