Skip to content

Commit b405435

Browse files
ptiurinptiurin
authored
feat(FIR-38127): Add struct support (#125)
1 parent 1a6cdaa commit b405435

File tree

5 files changed

+295
-2
lines changed

5 files changed

+295
-2
lines changed

‎src/statement/dataTypes.ts‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,8 @@ const getMappedType = (innerType: string) => {
4747
}
4848
};
4949

50-
const COMPLEX_TYPE = /(nullable|array)\((.+)\)/;
50+
const COMPLEX_TYPE = /^(nullable|array)\((.+)\)/;
51+
const STRUCT_TYPE = /^(struct)\((.+)\)/;
5152

5253
const DATE_TYPES = withNullableTypes([
5354
"pg_date",
@@ -97,6 +98,42 @@ export const getFireboltType = (type: string): string => {
9798
return mappedType || key;
9899
};
99100

101+
const trimElement = (element: string) =>
102+
// Remove leading and trailing spaces and backticks
103+
element.trim().replace(/(^\s*`?)|(`?\s*$)/g, "");
104+
105+
const decomposeSingleStructType = (type: string): [string, string] => {
106+
// Given a single struct element like "a int", extract the field and type
107+
// Finds the second backtick if any or the first space to separate field and type
108+
let index = type.indexOf("`", 1);
109+
if (index === -1) {
110+
index = type.indexOf(" ");
111+
}
112+
index = index + 1;
113+
const key = trimElement(type.substring(0, index));
114+
const value = trimElement(type.substring(index));
115+
return [key, value];
116+
};
117+
118+
export const getStructTypes = (type: string): Record<string, string> => {
119+
// Get a map of top level struct fields and their types, no recursion here
120+
// Example: "struct(a int, b struct(c text))" => { a: "int", b: "struct(c text)" }
121+
const match = STRUCT_TYPE.exec(type);
122+
if (match) {
123+
// extract types within struct
124+
const [_, _outerType, innerType] = match;
125+
// split types by comma (taking into account nested structs)
126+
const innerTypes = innerType.split(/,(?![^()]*\))/);
127+
const structTypes: Record<string, string> = {};
128+
for (const innerType of innerTypes) {
129+
const [field, type] = decomposeSingleStructType(innerType.trim());
130+
structTypes[field] = type;
131+
}
132+
return structTypes;
133+
}
134+
return {};
135+
};
136+
100137
export const getInnerType = (type: string): string => {
101138
const key = type.toLowerCase();
102139
const match = key.match(COMPLEX_TYPE);
@@ -127,3 +164,7 @@ export const isFloatType = (type: string) => {
127164
export const isNumberType = (type: string) => {
128165
return INTEGER_TYPES.includes(type) || isFloatType(type);
129166
};
167+
168+
export const isStructType = (type: string) => {
169+
return STRUCT_TYPE.test(type);
170+
};

‎src/statement/hydrateResponse.ts‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@ import {
55
isByteAType,
66
isDateType,
77
isNumberType,
8-
getInnerType
8+
getInnerType,
9+
getStructTypes,
10+
isStructType
911
} from "./dataTypes";
1012
import { hydrateDate } from "./hydrateDate";
1113

@@ -21,11 +23,39 @@ const hydrateInfNan = (value: string) => {
2123
return NaN;
2224
};
2325

26+
const hydrateStruct = (
27+
value: Record<string, unknown>,
28+
type: string,
29+
executeQueryOptions: ExecuteQueryOptions
30+
): Record<string, unknown> => {
31+
const hydratedStruct: Record<string, unknown> = {};
32+
const innerTypes = getStructTypes(type);
33+
// if number of keys does not match, return value as is
34+
if (Object.keys(innerTypes).length !== Object.keys(value).length) {
35+
return value;
36+
}
37+
for (const [key, innerType] of Object.entries(innerTypes)) {
38+
hydratedStruct[key] = getHydratedValue(
39+
value[key],
40+
innerType,
41+
executeQueryOptions
42+
);
43+
}
44+
return hydratedStruct;
45+
};
46+
2447
const getHydratedValue = (
2548
value: unknown,
2649
type: string,
2750
executeQueryOptions: ExecuteQueryOptions
2851
): any => {
52+
if (isStructType(type)) {
53+
return hydrateStruct(
54+
value as Record<string, unknown>,
55+
type,
56+
executeQueryOptions
57+
);
58+
}
2959
if (Array.isArray(value)) {
3060
const innerType = getInnerType(type);
3161
return value.map(element =>

‎test/integration/v2/fetchTypes.test.ts‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,4 +59,84 @@ describe("test type casting on fetch", () => {
5959
new BigNumber("-9223372036854775808")
6060
);
6161
});
62+
it("select geography", async () => {
63+
const firebolt = Firebolt({
64+
apiEndpoint: process.env.FIREBOLT_API_ENDPOINT as string
65+
});
66+
67+
const connection = await firebolt.connect(connectionParams);
68+
69+
const statement = await connection.execute(
70+
"select 'POINT(1 2)'::geography"
71+
);
72+
73+
const { data, meta } = await statement.fetchResult();
74+
expect(meta[0].type).toEqual("geography");
75+
const row = data[0];
76+
expect((row as unknown[])[0]).toEqual(
77+
"0101000020E6100000FEFFFFFFFFFFEF3F0000000000000040"
78+
);
79+
});
80+
it("select struct", async () => {
81+
const firebolt = Firebolt({
82+
apiEndpoint: process.env.FIREBOLT_API_ENDPOINT as string
83+
});
84+
85+
const connection = await firebolt.connect({
86+
...connectionParams,
87+
engineName: process.env.FIREBOLT_ENGINE_NAME as string
88+
});
89+
await connection.execute("SET advanced_mode=1");
90+
await connection.execute("SET enable_struct=1");
91+
await connection.execute("SET enable_create_table_v2=true");
92+
await connection.execute("SET enable_row_selection=true");
93+
await connection.execute("SET prevent_create_on_information_schema=true");
94+
await connection.execute("SET enable_create_table_with_struct_type=true");
95+
await connection.execute("DROP TABLE IF EXISTS test_struct");
96+
await connection.execute("DROP TABLE IF EXISTS test_struct_helper");
97+
try {
98+
await connection.execute(
99+
"CREATE TABLE IF NOT EXISTS test_struct(id int not null, s struct(a array(int) not null, b bytea null) not null)"
100+
);
101+
await connection.execute(
102+
"CREATE TABLE IF NOT EXISTS test_struct_helper(a array(int) not null, b bytea null)"
103+
);
104+
const bytea_value = Buffer.from("hello_world_123ツ\n\u0048");
105+
await connection.execute(
106+
"INSERT INTO test_struct_helper(a, b) VALUES ([1, 2], ?::bytea)",
107+
{ parameters: [bytea_value] }
108+
);
109+
// Test null values too
110+
await connection.execute(
111+
"INSERT INTO test_struct_helper(a, b) VALUES ([3, null], null)"
112+
);
113+
await connection.execute(
114+
"INSERT INTO test_struct(id, s) SELECT 1, test_struct_helper FROM test_struct_helper"
115+
);
116+
117+
const statement = await connection.execute(
118+
"SELECT test_struct FROM test_struct"
119+
);
120+
121+
const { data, meta } = await statement.fetchResult();
122+
expect(meta[0].type).toEqual(
123+
"struct(id int, s struct(a array(int null), b bytea null))"
124+
);
125+
const row = data[0];
126+
expect((row as unknown[])[0]).toEqual({
127+
id: 1,
128+
s: { a: [1, 2], b: bytea_value }
129+
});
130+
131+
const row2 = data[1];
132+
expect((row2 as unknown[])[0]).toEqual({
133+
id: 1,
134+
s: { a: [3, null], b: null }
135+
});
136+
} finally {
137+
// Make sure to always clean up
138+
await connection.execute("DROP TABLE IF EXISTS test_struct");
139+
await connection.execute("DROP TABLE IF EXISTS test_struct_helper");
140+
}
141+
});
62142
});

‎test/unit/hydrate.test.ts‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
import { getInnerType } from "../../src/statement/dataTypes";
2+
import { getStructTypes } from "../../src/statement/dataTypes";
3+
4+
describe("getInnerType function", () => {
5+
it("should return the inner type for a nullable type", () => {
6+
const type = "nullable(int)";
7+
const result = getInnerType(type);
8+
expect(result).toBe("int");
9+
});
10+
11+
it("should return the inner type for an array type", () => {
12+
const type = "array(int)";
13+
const result = getInnerType(type);
14+
expect(result).toBe("int");
15+
});
16+
17+
it("should return the inner type for a nested nullable type", () => {
18+
const type = "nullable(nullable(int))";
19+
const result = getInnerType(type);
20+
expect(result).toBe("int");
21+
});
22+
23+
it("should return the original type if no complex type is found", () => {
24+
const type = "int";
25+
const result = getInnerType(type);
26+
expect(result).toBe("int");
27+
});
28+
});
29+
30+
describe("getStructTypes function", () => {
31+
it("should return the correct types for a simple struct", () => {
32+
const type = "struct(a int, b text)";
33+
const result = getStructTypes(type);
34+
expect(result).toEqual({ a: "int", b: "text" });
35+
});
36+
37+
it("should return the correct types for a nested struct", () => {
38+
const type = "struct(a int, b struct(c int, d text))";
39+
const result = getStructTypes(type);
40+
expect(result).toEqual({ a: "int", b: "struct(c int, d text)" });
41+
});
42+
43+
it("should return an empty object for a non-struct type", () => {
44+
const type = "int";
45+
const result = getStructTypes(type);
46+
expect(result).toEqual({});
47+
});
48+
49+
it("should handle structs with nullable types", () => {
50+
const type = "struct(a int null, b text null)";
51+
const result = getStructTypes(type);
52+
expect(result).toEqual({ a: "int null", b: "text null" });
53+
});
54+
55+
it("should handle structs with array types", () => {
56+
const type = "struct(a array(int), b array(text))";
57+
const result = getStructTypes(type);
58+
expect(result).toEqual({ a: "array(int)", b: "array(text)" });
59+
});
60+
61+
it("should handle structs with nested array types", () => {
62+
const type = "struct(a array(int), b array(array(text)))";
63+
const result = getStructTypes(type);
64+
expect(result).toEqual({ a: "array(int)", b: "array(array(text))" });
65+
});
66+
67+
it("should not fail on malformed struct types", () => {
68+
const type = "struct(a int, b text";
69+
const result = getStructTypes(type);
70+
// This allows the outer function to continue processing the type
71+
// as text
72+
expect(result).toEqual({});
73+
});
74+
75+
it("should handle structs with mixed case columns", () => {
76+
const type = "struct(mIxEdCaSe int, MiXeDcAsE text)";
77+
const result = getStructTypes(type);
78+
expect(result).toEqual({ mIxEdCaSe: "int", MiXeDcAsE: "text" });
79+
});
80+
81+
it("should handle structs with spaces in column names", () => {
82+
const type = "struct(`column name` int, `column name 2` text)";
83+
const result = getStructTypes(type);
84+
expect(result).toEqual({ "column name": "int", "column name 2": "text" });
85+
});
86+
87+
it("should respect spaces within backticks", () => {
88+
const type = "struct(` column name ` int, ` column name 2 ` text)";
89+
const result = getStructTypes(type);
90+
expect(result).toEqual({
91+
" column name ": "int",
92+
" column name 2 ": "text"
93+
});
94+
});
95+
});

‎test/unit/v2/statement.test.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,53 @@ describe("parse values", () => {
313313
new BigNumber(1000000000000000000000000000000000000)
314314
);
315315
});
316+
it("parses bytea into Buffer", () => {
317+
const row = {
318+
bytea: "\\x68656c6c6f5f776f726c64"
319+
};
320+
const meta = [{ name: "bytea", type: "bytea" }];
321+
const res: Record<string, Buffer> = hydrateRow(row, meta, {});
322+
expect(res["bytea"]).toEqual(Buffer.from("hello_world"));
323+
});
324+
it("parses struct into object", () => {
325+
const row = {
326+
s: { a: [1, 2], b: "\\x68656c6c6f5f776f726c64" }
327+
};
328+
const meta = [
329+
{ name: "s", type: "struct(a array(int null), b bytea null)" }
330+
];
331+
const res: Record<string, Record<string, any>> = hydrateRow(row, meta, {});
332+
expect(res["s"]).toEqual({ a: [1, 2], b: Buffer.from("hello_world") });
333+
});
334+
it("parses nested struct into object", () => {
335+
const row = {
336+
s: { a: [1, 2], b: { c: "hello" } }
337+
};
338+
const meta = [
339+
{ name: "s", type: "struct(a array(int null), b struct(c text))" }
340+
];
341+
const res: Record<string, Record<string, any>> = hydrateRow(row, meta, {});
342+
expect(res["s"]).toEqual({ a: [1, 2], b: { c: "hello" } });
343+
});
344+
it("parses nested struct with nulls correctly", () => {
345+
const row = {
346+
s: { a: [1, null], b: { c: null } }
347+
};
348+
const meta = [
349+
{ name: "s", type: "struct(a array(int null), b struct(c text null))" }
350+
];
351+
const res: Record<string, Record<string, any>> = hydrateRow(row, meta, {});
352+
expect(res["s"]).toEqual({ a: [1, null], b: { c: null } });
353+
});
354+
it("does not break on malformed struct", () => {
355+
const row = {
356+
s: { a: [1, 2], b: "hello" }
357+
};
358+
// Missing closing parenthesis
359+
const meta = [{ name: "s", type: "struct(a array(int), b text" }];
360+
const res: Record<string, Record<string, any>> = hydrateRow(row, meta, {});
361+
expect(res["s"]).toEqual({ a: [1, 2], b: "hello" });
362+
});
316363
});
317364

318365
describe("set statements", () => {

0 commit comments

Comments
 (0)