From f5141353c17c6f8cf9118ddb8ea7c87278630aac Mon Sep 17 00:00:00 2001
From: Kyle Gray <kyle@sqlc.dev>
Date: Wed, 18 Oct 2023 09:42:23 -0700
Subject: [PATCH] feat(postgresql): Support system columns on tables (#2871)

* test: Add a test case for https://github.com/sqlc-dev/sqlc/issues/1745

* feat(postgresql): Support system columns

* Remove stdlib

* Update sqlc.yaml

* Add exec.json

---------

Co-authored-by: Andrew Benton <andrew@sqlc.dev>
---
 .../endtoend/testdata/select_system/issue.md  |  1 +
 .../testdata/select_system/pgx/exec.json      |  3 ++
 .../testdata/select_system/pgx/go/db.go       | 32 +++++++++++++++
 .../testdata/select_system/pgx/go/models.go   | 13 ++++++
 .../select_system/pgx/go/query.sql.go         | 41 +++++++++++++++++++
 .../testdata/select_system/pgx/query.sql      |  4 ++
 .../testdata/select_system/pgx/schema.sql     |  3 ++
 .../testdata/select_system/pgx/sqlc.yaml      | 10 +++++
 .../engine/postgresql/analyzer/analyze.go     |  2 +-
 9 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 internal/endtoend/testdata/select_system/issue.md
 create mode 100644 internal/endtoend/testdata/select_system/pgx/exec.json
 create mode 100644 internal/endtoend/testdata/select_system/pgx/go/db.go
 create mode 100644 internal/endtoend/testdata/select_system/pgx/go/models.go
 create mode 100644 internal/endtoend/testdata/select_system/pgx/go/query.sql.go
 create mode 100644 internal/endtoend/testdata/select_system/pgx/query.sql
 create mode 100644 internal/endtoend/testdata/select_system/pgx/schema.sql
 create mode 100644 internal/endtoend/testdata/select_system/pgx/sqlc.yaml

diff --git a/internal/endtoend/testdata/select_system/issue.md b/internal/endtoend/testdata/select_system/issue.md
new file mode 100644
index 0000000000..7e536076fc
--- /dev/null
+++ b/internal/endtoend/testdata/select_system/issue.md
@@ -0,0 +1 @@
+https://github.com/sqlc-dev/sqlc/issues/1745
diff --git a/internal/endtoend/testdata/select_system/pgx/exec.json b/internal/endtoend/testdata/select_system/pgx/exec.json
new file mode 100644
index 0000000000..e5dfda7818
--- /dev/null
+++ b/internal/endtoend/testdata/select_system/pgx/exec.json
@@ -0,0 +1,3 @@
+{
+    "contexts": ["managed-db"]
+}
diff --git a/internal/endtoend/testdata/select_system/pgx/go/db.go b/internal/endtoend/testdata/select_system/pgx/go/db.go
new file mode 100644
index 0000000000..8a010ccc48
--- /dev/null
+++ b/internal/endtoend/testdata/select_system/pgx/go/db.go
@@ -0,0 +1,32 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.22.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/select_system/pgx/go/models.go b/internal/endtoend/testdata/select_system/pgx/go/models.go
new file mode 100644
index 0000000000..b52387cf1a
--- /dev/null
+++ b/internal/endtoend/testdata/select_system/pgx/go/models.go
@@ -0,0 +1,13 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.22.0
+
+package querytest
+
+import (
+	"github.com/jackc/pgx/v5/pgtype"
+)
+
+type Test struct {
+	ID pgtype.Int4
+}
diff --git a/internal/endtoend/testdata/select_system/pgx/go/query.sql.go b/internal/endtoend/testdata/select_system/pgx/go/query.sql.go
new file mode 100644
index 0000000000..fadb461a63
--- /dev/null
+++ b/internal/endtoend/testdata/select_system/pgx/go/query.sql.go
@@ -0,0 +1,41 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.22.0
+// source: query.sql
+
+package querytest
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5/pgtype"
+)
+
+const getSystemColumns = `-- name: GetSystemColumns :one
+SELECT
+  tableoid, xmin, cmin, xmax, cmax, ctid
+FROM test
+`
+
+type GetSystemColumnsRow struct {
+	Tableoid pgtype.Uint32
+	Xmin     pgtype.Uint32
+	Cmin     pgtype.Uint32
+	Xmax     pgtype.Uint32
+	Cmax     pgtype.Uint32
+	Ctid     pgtype.TID
+}
+
+func (q *Queries) GetSystemColumns(ctx context.Context) (GetSystemColumnsRow, error) {
+	row := q.db.QueryRow(ctx, getSystemColumns)
+	var i GetSystemColumnsRow
+	err := row.Scan(
+		&i.Tableoid,
+		&i.Xmin,
+		&i.Cmin,
+		&i.Xmax,
+		&i.Cmax,
+		&i.Ctid,
+	)
+	return i, err
+}
diff --git a/internal/endtoend/testdata/select_system/pgx/query.sql b/internal/endtoend/testdata/select_system/pgx/query.sql
new file mode 100644
index 0000000000..11ddeb7495
--- /dev/null
+++ b/internal/endtoend/testdata/select_system/pgx/query.sql
@@ -0,0 +1,4 @@
+-- name: GetSystemColumns :one
+SELECT
+  tableoid, xmin, cmin, xmax, cmax, ctid
+FROM test;
diff --git a/internal/endtoend/testdata/select_system/pgx/schema.sql b/internal/endtoend/testdata/select_system/pgx/schema.sql
new file mode 100644
index 0000000000..7518997c95
--- /dev/null
+++ b/internal/endtoend/testdata/select_system/pgx/schema.sql
@@ -0,0 +1,3 @@
+CREATE TABLE test (
+  id INT
+);
diff --git a/internal/endtoend/testdata/select_system/pgx/sqlc.yaml b/internal/endtoend/testdata/select_system/pgx/sqlc.yaml
new file mode 100644
index 0000000000..5dc63e3f91
--- /dev/null
+++ b/internal/endtoend/testdata/select_system/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"
diff --git a/internal/engine/postgresql/analyzer/analyze.go b/internal/engine/postgresql/analyzer/analyze.go
index ab34eb348e..fe07b840e6 100644
--- a/internal/engine/postgresql/analyzer/analyze.go
+++ b/internal/engine/postgresql/analyzer/analyze.go
@@ -101,7 +101,7 @@ func (a *Analyzer) columnInfo(ctx context.Context, field pgconn.FieldDescription
 	if ok {
 		return cinfo.(*pgColumn), nil
 	}
-	rows, err := a.pool.Query(ctx, columnQuery, field.TableOID, field.TableAttributeNumber)
+	rows, err := a.pool.Query(ctx, columnQuery, field.TableOID, int16(field.TableAttributeNumber))
 	if err != nil {
 		return nil, err
 	}