Skip to content

Commit dd478cd

Browse files
exAspArkexAspArk
committed
Release v2.0.1
1 parent d6e4a15 commit dd478cd

4 files changed

Lines changed: 175 additions & 14 deletions

File tree

‎server/config.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import (
99
)
1010

1111
const (
12-
VERSION = "2.0.0"
12+
VERSION = "2.0.1"
1313

1414
ENV_LOG_LEVEL = "LOG_LEVEL"
1515
ENV_AUDIT_ENABLED = "AUDIT_ENABLED"

‎server/query_handler_test.go‎

Lines changed: 117 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,113 @@ type queryResponseExpectation struct {
2020
values []string
2121
}
2222

23+
func TestHandleQueryExpandArray(t *testing.T) {
24+
queryHandler := initQueryHandlerForTest(t)
25+
tests := []struct {
26+
query string
27+
columns []string
28+
rows [][]string
29+
}{
30+
{
31+
query: "SELECT information_schema._pg_expandarray(ARRAY[1,2])",
32+
columns: []string{"_pg_expandarray"},
33+
rows: [][]string{{"(1,1)"}, {"(2,2)"}},
34+
},
35+
{
36+
query: "SELECT (information_schema._pg_expandarray(ARRAY[10,NULL,30])).x AS x, (information_schema._pg_expandarray(ARRAY[10,NULL,30])).n AS n",
37+
columns: []string{"x", "n"},
38+
rows: [][]string{{"10", "1"}, {"", "2"}, {"30", "3"}},
39+
},
40+
{
41+
query: "SELECT * FROM information_schema._pg_expandarray(ARRAY[10,20])",
42+
columns: []string{"x", "n"},
43+
rows: [][]string{{"10", "1"}, {"20", "2"}},
44+
},
45+
{
46+
query: "SELECT information_schema._pg_expandarray(ARRAY[]::integer[])",
47+
columns: []string{"_pg_expandarray"},
48+
},
49+
{
50+
query: "SELECT information_schema._pg_expandarray(NULL::integer[])",
51+
columns: []string{"_pg_expandarray"},
52+
},
53+
{
54+
query: "SELECT information_schema._pg_expandarray(ARRAY[NULL,1])",
55+
columns: []string{"_pg_expandarray"},
56+
rows: [][]string{{"(,1)"}, {"(1,2)"}},
57+
},
58+
{
59+
query: `SELECT information_schema._pg_expandarray(ARRAY['', 'a,b', 'a"b', 'a\b'])`,
60+
columns: []string{"_pg_expandarray"},
61+
rows: [][]string{{`("",1)`}, {`("a,b",2)`}, {`("a""b",3)`}, {`("a\\b",4)`}},
62+
},
63+
{
64+
// JDBC expands the same array twice and accesses the composite in an outer query.
65+
query: `SELECT result.key_seq, (result.keys).x AS column_number FROM (
66+
SELECT (information_schema._pg_expandarray(i.indkey)).n AS key_seq,
67+
information_schema._pg_expandarray(i.indkey) AS keys
68+
FROM (SELECT ARRAY[3,1]::smallint[] AS indkey) i
69+
) result ORDER BY result.key_seq`,
70+
columns: []string{"key_seq", "column_number"},
71+
rows: [][]string{{"1", "3"}, {"2", "1"}},
72+
},
73+
}
74+
for _, test := range tests {
75+
t.Run(test.query, func(t *testing.T) {
76+
messages, err := queryHandler.HandleSimpleQuery(test.query)
77+
if err != nil {
78+
t.Fatal(err)
79+
}
80+
if len(messages) != len(test.rows)+2 {
81+
t.Fatalf("message count = %d, want %d", len(messages), len(test.rows)+2)
82+
}
83+
description := messages[0].(*pgproto3.RowDescription)
84+
columns := make([]string, len(description.Fields))
85+
for i, field := range description.Fields {
86+
columns[i] = string(field.Name)
87+
}
88+
if !slices.Equal(columns, test.columns) {
89+
t.Fatalf("columns = %v, want %v", columns, test.columns)
90+
}
91+
if test.columns[0] == "_pg_expandarray" && description.Fields[0].DataTypeOID != pgtype.RecordOID {
92+
t.Fatalf("record OID = %d, want %d", description.Fields[0].DataTypeOID, pgtype.RecordOID)
93+
}
94+
for i, row := range test.rows {
95+
assertDataRow(t, messages[i+1], row)
96+
}
97+
})
98+
}
99+
}
100+
101+
func TestHandleQueryJDBCPrimaryKeysWithoutPrimaryKey(t *testing.T) {
102+
queryHandler := initQueryHandlerForTest(t)
103+
if _, err := queryHandler.ServerDuckdbClient.ExecContext(context.Background(), "CREATE SCHEMA audit; CREATE TABLE audit.changes(id INTEGER)"); err != nil {
104+
t.Fatal(err)
105+
}
106+
// pgjdbc getPrimaryKeys query, including the composite field access used by Airbyte.
107+
assertQueryResponse(t, queryHandler, queryResponseExpectation{
108+
query: `SELECT result.TABLE_CAT AS "TABLE_CAT", result.TABLE_SCHEM AS "TABLE_SCHEM",
109+
result.TABLE_NAME AS "TABLE_NAME", result.COLUMN_NAME AS "COLUMN_NAME",
110+
result.KEY_SEQ AS "KEY_SEQ", result.PK_NAME AS "PK_NAME"
111+
FROM (
112+
SELECT current_database() AS TABLE_CAT, n.nspname AS TABLE_SCHEM,
113+
ct.relname AS TABLE_NAME, a.attname AS COLUMN_NAME,
114+
(information_schema._pg_expandarray(i.indkey)).n AS KEY_SEQ, ci.relname AS PK_NAME,
115+
information_schema._pg_expandarray(i.indkey) AS KEYS, a.attnum AS A_ATTNUM,
116+
i.indnkeyatts AS NUM_KEYS
117+
FROM pg_catalog.pg_class ct
118+
JOIN pg_catalog.pg_attribute a ON (ct.oid = a.attrelid)
119+
JOIN pg_catalog.pg_namespace n ON (ct.relnamespace = n.oid)
120+
JOIN pg_catalog.pg_index i ON (a.attrelid = i.indrelid)
121+
JOIN pg_catalog.pg_class ci ON (ci.oid = i.indexrelid)
122+
WHERE true AND n.nspname = 'audit' AND ct.relname = 'changes' AND i.indisprimary
123+
) result
124+
WHERE result.A_ATTNUM = (result.KEYS).x AND result.KEY_SEQ <= result.NUM_KEYS
125+
ORDER BY result.table_name, result.pk_name, result.key_seq`,
126+
columnNames: []string{"TABLE_CAT", "TABLE_SCHEM", "TABLE_NAME", "COLUMN_NAME", "KEY_SEQ", "PK_NAME"},
127+
})
128+
}
129+
23130
func TestHandleQueryFixedSizeFloatArray(t *testing.T) {
24131
queryHandler := initQueryHandlerForTest(t)
25132
if _, err := queryHandler.ServerDuckdbClient.ExecContext(
@@ -488,18 +595,16 @@ func TestHandleQueryExpressions(t *testing.T) {
488595
columnNames: []string{"index"},
489596
values: []string{"1"},
490597
},
491-
// Temporarily disabled: _pg_expandarray is registered as a DuckDB table
492-
// function, but PostgreSQL permits this composite-returning scalar syntax.
493-
// {
494-
// query: "SELECT (information_schema._pg_expandarray(ARRAY[10])).n",
495-
// columnNames: []string{"n"},
496-
// values: []string{"1"},
497-
// },
498-
// {
499-
// query: "SELECT (information_schema._pg_expandarray(ARRAY[10])).x AS value",
500-
// columnNames: []string{"value"},
501-
// values: []string{"10"},
502-
// },
598+
{
599+
query: "SELECT (information_schema._pg_expandarray(ARRAY[10])).n",
600+
columnNames: []string{"n"},
601+
values: []string{"1"},
602+
},
603+
{
604+
query: "SELECT (information_schema._pg_expandarray(ARRAY[10])).x AS value",
605+
columnNames: []string{"value"},
606+
values: []string{"10"},
607+
},
503608
{
504609
query: "SELECT s.usename, r.rolconfig FROM pg_catalog.pg_shadow s LEFT JOIN pg_catalog.pg_roles r ON s.usename = r.rolname",
505610
columnNames: []string{"usename", "rolconfig"},

‎server/query_remapper_function.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,10 @@ func CreatePgCatalogMacroQueries(config *Config) []string {
103103

104104
func CreateInformationSchemaMacroQueries(config *Config) []string {
105105
result := []string{
106-
"CREATE MACRO _pg_expandarray(arr) AS TABLE SELECT unnest(arr) AS x, unnest(generate_series(1, array_length(arr))) AS n",
106+
// PostgreSQL permits this set-returning function in both SELECT and FROM.
107+
// Keep x and its one-based position together so composite field access works.
108+
"CREATE MACRO _pg_expandarray(arr) AS unnest(list_transform(arr, lambda value, position: struct_pack(x := value, n := position::INTEGER)))",
109+
"CREATE MACRO _pg_expandarray(arr) AS TABLE SELECT unnest(arr) AS x, unnest(generate_series(1, array_length(arr)))::INTEGER AS n",
107110
}
108111
PG_INFORMATION_SCHEMA_MACRO_FUNCTION_NAMES = extractMacroNames(result)
109112
return result

‎server/response_handler.go‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,12 +113,18 @@ func (responseHandler *ResponseHandler) ColumnDescriptionTypeOid(col *sql.Column
113113
return pgtype.NumericOID
114114
}
115115
}
116+
if isExpandedArrayRecord(databaseTypeName) {
117+
return pgtype.RecordOID
118+
}
116119

117120
Panic(responseHandler.Config, "Unsupported serialized column type: "+databaseTypeName)
118121
return 0
119122
}
120123

121124
func (responseHandler *ResponseHandler) RowValuePointer(col *sql.ColumnType) interface{} {
125+
if isExpandedArrayRecord(col.DatabaseTypeName()) {
126+
return new(NullExpandedArrayRecord)
127+
}
122128
if col.DatabaseTypeName() == "UINTEGER" || col.DatabaseTypeName() == "UBIGINT" {
123129
return new(NullUint64)
124130
}
@@ -239,6 +245,11 @@ func (responseHandler *ResponseHandler) RowValueBytes(valuePtr interface{}, col
239245
} else {
240246
return nil
241247
}
248+
case *NullExpandedArrayRecord:
249+
if value.Present {
250+
return []byte(value.String())
251+
}
252+
return nil
242253
case *string:
243254
return []byte(*value)
244255
}
@@ -247,6 +258,48 @@ func (responseHandler *ResponseHandler) RowValueBytes(valuePtr interface{}, col
247258
return nil
248259
}
249260

261+
// _pg_expandarray returns a composite with fields in the order (x, n).
262+
func isExpandedArrayRecord(databaseTypeName string) bool {
263+
return strings.HasPrefix(databaseTypeName, `STRUCT("x" `) && strings.HasSuffix(databaseTypeName, `, "n" INTEGER)`)
264+
}
265+
266+
type NullExpandedArrayRecord struct {
267+
Present bool
268+
Value map[string]interface{}
269+
}
270+
271+
func (record *NullExpandedArrayRecord) Scan(value interface{}) error {
272+
record.Present = value != nil
273+
if value == nil {
274+
return nil
275+
}
276+
fields, ok := value.(map[string]interface{})
277+
if !ok {
278+
return fmt.Errorf("unexpected _pg_expandarray record type: %T", value)
279+
}
280+
record.Value = fields
281+
return nil
282+
}
283+
284+
func (record NullExpandedArrayRecord) String() string {
285+
fields := make([]string, 2)
286+
for i, name := range []string{"x", "n"} {
287+
value := record.Value[name]
288+
if value == nil {
289+
continue
290+
}
291+
field := fmt.Sprint(value)
292+
if boolean, ok := value.(bool); ok {
293+
field = strconv.FormatBool(boolean)[:1]
294+
}
295+
if field == "" || strings.ContainsAny(field, ",()\"\\ \t\n\r\v\f") {
296+
field = `"` + strings.NewReplacer(`\`, `\\`, `"`, `""`).Replace(field) + `"`
297+
}
298+
fields[i] = field
299+
}
300+
return "(" + strings.Join(fields, ",") + ")"
301+
}
302+
250303
func (responseHandler *ResponseHandler) isSystemTableOidColumn(colName string) bool {
251304
oidColumns := map[string]bool{
252305
"oid": true,

0 commit comments

Comments
 (0)