Skip to content

Commit 59f2958

Browse files
dankochetovdankochetov
authored
Merge pull request #1911 from livingforjesus/fix-arrayvalue-issue
[Pg] Add insert/update array support in aws-data-api
2 parents 2c9b73b + 338650d commit 59f2958

File tree

3 files changed

+171
-14
lines changed

3 files changed

+171
-14
lines changed

‎drizzle-orm/src/aws-data-api/common/index.ts

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,19 @@ export function getValueFromDataApi(field: Field) {
2020
if (field.arrayValue.stringValues !== undefined) {
2121
return field.arrayValue.stringValues;
2222
}
23+
if (field.arrayValue.longValues !== undefined) {
24+
return field.arrayValue.longValues;
25+
}
26+
if (field.arrayValue.doubleValues !== undefined) {
27+
return field.arrayValue.doubleValues;
28+
}
29+
if (field.arrayValue.booleanValues !== undefined) {
30+
return field.arrayValue.booleanValues;
31+
}
32+
if (field.arrayValue.arrayValues !== undefined) {
33+
return field.arrayValue.arrayValues;
34+
}
35+
2336
throw new Error('Unknown array type');
2437
} else {
2538
throw new Error('Unknown type');

‎drizzle-orm/src/aws-data-api/pg/driver.ts

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,20 @@
1-
import { entityKind } from '~/entity.ts';
2-
import type { SQLWrapper } from '~/index.ts';
1+
import { entityKind, is } from '~/entity.ts';
2+
import type { SQL, SQLWrapper } from '~/index.ts';
3+
import { Param, sql, Table } from '~/index.ts';
34
import type { Logger } from '~/logger.ts';
45
import { DefaultLogger } from '~/logger.ts';
56
import { PgDatabase } from '~/pg-core/db.ts';
67
import { PgDialect } from '~/pg-core/dialect.ts';
8+
import type { PgColumn, PgInsertConfig, PgTable, TableConfig } from '~/pg-core/index.ts';
9+
import { PgArray } from '~/pg-core/index.ts';
710
import type { PgRaw } from '~/pg-core/query-builders/raw.ts';
811
import {
912
createTableRelationsHelpers,
1013
extractTablesRelationalConfig,
1114
type RelationalSchemaConfig,
1215
type TablesRelationalConfig,
1316
} from '~/relations.ts';
14-
import type { DrizzleConfig } from '~/utils.ts';
17+
import type { DrizzleConfig, UpdateSet } from '~/utils.ts';
1518
import type { AwsDataApiClient, AwsDataApiPgQueryResult, AwsDataApiPgQueryResultHKT } from './session.ts';
1619
import { AwsDataApiSession } from './session.ts';
1720

@@ -48,6 +51,45 @@ export class AwsPgDialect extends PgDialect {
4851
override escapeParam(num: number): string {
4952
return `:${num + 1}`;
5053
}
54+
55+
override buildInsertQuery(
56+
{ table, values, onConflict, returning }: PgInsertConfig<PgTable<TableConfig>>,
57+
): SQL<unknown> {
58+
const columns: Record<string, PgColumn> = table[Table.Symbol.Columns];
59+
const colEntries: [string, PgColumn][] = Object.entries(columns);
60+
for (const value of values) {
61+
for (const [fieldName, col] of colEntries) {
62+
const colValue = value[fieldName];
63+
if (
64+
is(colValue, Param) && colValue.value !== undefined && is(colValue.encoder, PgArray)
65+
&& Array.isArray(colValue.value)
66+
) {
67+
value[fieldName] = sql`cast(${col.mapToDriverValue(colValue.value)} as ${
68+
sql.raw(colValue.encoder.getSQLType())
69+
})`;
70+
}
71+
}
72+
}
73+
74+
return super.buildInsertQuery({ table, values, onConflict, returning });
75+
}
76+
77+
override buildUpdateSet(table: PgTable<TableConfig>, set: UpdateSet): SQL<unknown> {
78+
const columns: Record<string, PgColumn> = table[Table.Symbol.Columns];
79+
80+
for (const [colName, colValue] of Object.entries(set)) {
81+
const currentColumn = columns[colName];
82+
if (
83+
currentColumn && is(colValue, Param) && colValue.value !== undefined && is(colValue.encoder, PgArray)
84+
&& Array.isArray(colValue.value)
85+
) {
86+
set[colName] = sql`cast(${currentColumn?.mapToDriverValue(colValue.value)} as ${
87+
sql.raw(colValue.encoder.getSQLType())
88+
})`;
89+
}
90+
}
91+
return super.buildUpdateSet(table, set);
92+
}
5193
}
5294

5395
export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(

‎integration-tests/tests/awsdatapi.test.ts

Lines changed: 113 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ const usersTable = pgTable('users', {
3434
name: text('name').notNull(),
3535
verified: boolean('verified').notNull().default(false),
3636
jsonb: jsonb('jsonb').$type<string[]>(),
37+
bestTexts: text('best_texts').array().default(sql`'{}'`).notNull(),
3738
createdAt: timestamp('created_at', { withTimezone: true }).notNull().defaultNow(),
3839
});
3940

@@ -69,6 +70,7 @@ beforeEach(async () => {
6970
name text not null,
7071
verified boolean not null default false,
7172
jsonb jsonb,
73+
best_texts text[] not null default '{}',
7274
created_at timestamptz not null default now()
7375
)
7476
`,
@@ -84,7 +86,14 @@ test('select all fields', async () => {
8486

8587
expect(result[0]!.createdAt).toBeInstanceOf(Date);
8688
// t.assert(Math.abs(result[0]!.createdAt.getTime() - now) < 100);
87-
expect(result).toEqual([{ id: 1, name: 'John', verified: false, jsonb: null, createdAt: result[0]!.createdAt }]);
89+
expect(result).toEqual([{
90+
bestTexts: [],
91+
id: 1,
92+
name: 'John',
93+
verified: false,
94+
jsonb: null,
95+
createdAt: result[0]!.createdAt,
96+
}]);
8897
});
8998

9099
test('select sql', async () => {
@@ -176,7 +185,14 @@ test('update with returning all fields', async () => {
176185

177186
expect(users[0]!.createdAt).toBeInstanceOf(Date);
178187
// t.assert(Math.abs(users[0]!.createdAt.getTime() - now) < 100);
179-
expect(users).toEqual([{ id: 1, name: 'Jane', verified: false, jsonb: null, createdAt: users[0]!.createdAt }]);
188+
expect(users).toEqual([{
189+
id: 1,
190+
bestTexts: [],
191+
name: 'Jane',
192+
verified: false,
193+
jsonb: null,
194+
createdAt: users[0]!.createdAt,
195+
}]);
180196
});
181197

182198
test('update with returning partial', async () => {
@@ -195,7 +211,14 @@ test('delete with returning all fields', async () => {
195211

196212
expect(users[0]!.createdAt).toBeInstanceOf(Date);
197213
// t.assert(Math.abs(users[0]!.createdAt.getTime() - now) < 100);
198-
expect(users).toEqual([{ id: 1, name: 'John', verified: false, jsonb: null, createdAt: users[0]!.createdAt }]);
214+
expect(users).toEqual([{
215+
bestTexts: [],
216+
id: 1,
217+
name: 'John',
218+
verified: false,
219+
jsonb: null,
220+
createdAt: users[0]!.createdAt,
221+
}]);
199222
});
200223

201224
test('delete with returning partial', async () => {
@@ -211,13 +234,20 @@ test('delete with returning partial', async () => {
211234
test('insert + select', async () => {
212235
await db.insert(usersTable).values({ name: 'John' });
213236
const result = await db.select().from(usersTable);
214-
expect(result).toEqual([{ id: 1, name: 'John', verified: false, jsonb: null, createdAt: result[0]!.createdAt }]);
237+
expect(result).toEqual([{
238+
bestTexts: [],
239+
id: 1,
240+
name: 'John',
241+
verified: false,
242+
jsonb: null,
243+
createdAt: result[0]!.createdAt,
244+
}]);
215245

216246
await db.insert(usersTable).values({ name: 'Jane' });
217247
const result2 = await db.select().from(usersTable);
218248
expect(result2).toEqual([
219-
{ id: 1, name: 'John', verified: false, jsonb: null, createdAt: result2[0]!.createdAt },
220-
{ id: 2, name: 'Jane', verified: false, jsonb: null, createdAt: result2[1]!.createdAt },
249+
{ bestTexts: [], id: 1, name: 'John', verified: false, jsonb: null, createdAt: result2[0]!.createdAt },
250+
{ bestTexts: [], id: 2, name: 'Jane', verified: false, jsonb: null, createdAt: result2[1]!.createdAt },
221251
]);
222252
});
223253

@@ -236,7 +266,14 @@ test('insert with overridden default values', async () => {
236266
await db.insert(usersTable).values({ name: 'John', verified: true });
237267
const result = await db.select().from(usersTable);
238268

239-
expect(result).toEqual([{ id: 1, name: 'John', verified: true, jsonb: null, createdAt: result[0]!.createdAt }]);
269+
expect(result).toEqual([{
270+
bestTexts: [],
271+
id: 1,
272+
name: 'John',
273+
verified: true,
274+
jsonb: null,
275+
createdAt: result[0]!.createdAt,
276+
}]);
240277
});
241278

242279
test('insert many', async () => {
@@ -385,12 +422,14 @@ test('full join with alias', async () => {
385422
expect(result).toEqual([{
386423
users: {
387424
id: 10,
425+
bestTexts: [],
388426
name: 'Ivan',
389427
verified: false,
390428
jsonb: null,
391429
createdAt: result[0]!.users.createdAt,
392430
},
393431
customer: {
432+
bestTexts: [],
394433
id: 11,
395434
name: 'Hans',
396435
verified: false,
@@ -627,7 +666,7 @@ test('build query insert with onConflict do update', async () => {
627666

628667
expect(query).toEqual({
629668
sql:
630-
'insert into "users" ("id", "name", "verified", "jsonb", "created_at") values (default, :1, default, :2, default) on conflict ("id") do update set "name" = :3',
669+
'insert into "users" ("id", "name", "verified", "jsonb", "best_texts", "created_at") values (default, :1, default, :2, default, default) on conflict ("id") do update set "name" = :3',
631670
params: ['John', '["foo","bar"]', 'John1'],
632671
// typings: ['none', 'json', 'none']
633672
});
@@ -641,7 +680,7 @@ test('build query insert with onConflict do update / multiple columns', async ()
641680

642681
expect(query).toEqual({
643682
sql:
644-
'insert into "users" ("id", "name", "verified", "jsonb", "created_at") values (default, :1, default, :2, default) on conflict ("id","name") do update set "name" = :3',
683+
'insert into "users" ("id", "name", "verified", "jsonb", "best_texts", "created_at") values (default, :1, default, :2, default, default) on conflict ("id","name") do update set "name" = :3',
645684
params: ['John', '["foo","bar"]', 'John1'],
646685
// typings: ['none', 'json', 'none']
647686
});
@@ -655,7 +694,7 @@ test('build query insert with onConflict do nothing', async () => {
655694

656695
expect(query).toEqual({
657696
sql:
658-
'insert into "users" ("id", "name", "verified", "jsonb", "created_at") values (default, :1, default, :2, default) on conflict do nothing',
697+
'insert into "users" ("id", "name", "verified", "jsonb", "best_texts", "created_at") values (default, :1, default, :2, default, default) on conflict do nothing',
659698
params: ['John', '["foo","bar"]'],
660699
// typings: ['none', 'json']
661700
});
@@ -669,7 +708,7 @@ test('build query insert with onConflict do nothing + target', async () => {
669708

670709
expect(query).toEqual({
671710
sql:
672-
'insert into "users" ("id", "name", "verified", "jsonb", "created_at") values (default, :1, default, :2, default) on conflict ("id") do nothing',
711+
'insert into "users" ("id", "name", "verified", "jsonb", "best_texts", "created_at") values (default, :1, default, :2, default, default) on conflict ("id") do nothing',
673712
params: ['John', '["foo","bar"]'],
674713
// typings: ['none', 'json']
675714
});
@@ -857,6 +896,69 @@ test('select from raw sql with mapped values', async () => {
857896
]);
858897
});
859898

899+
test('insert with array values works', async () => {
900+
const bestTexts = ['text1', 'text2', 'text3'];
901+
const [insertResult] = await db.insert(usersTable).values({
902+
name: 'John',
903+
bestTexts,
904+
}).returning();
905+
906+
expect(insertResult?.bestTexts).toEqual(bestTexts);
907+
});
908+
909+
test('update with array values works', async () => {
910+
const [newUser] = await db.insert(usersTable).values({ name: 'John' }).returning();
911+
912+
const bestTexts = ['text4', 'text5', 'text6'];
913+
const [insertResult] = await db.update(usersTable).set({
914+
bestTexts,
915+
}).where(eq(usersTable.id, newUser!.id)).returning();
916+
917+
expect(insertResult?.bestTexts).toEqual(bestTexts);
918+
});
919+
920+
test('insert with array values works', async () => {
921+
const bestTexts = ['text1', 'text2', 'text3'];
922+
const [insertResult] = await db.insert(usersTable).values({
923+
name: 'John',
924+
bestTexts,
925+
}).returning();
926+
927+
expect(insertResult?.bestTexts).toEqual(bestTexts);
928+
});
929+
930+
test('update with array values works', async () => {
931+
const [newUser] = await db.insert(usersTable).values({ name: 'John' }).returning();
932+
933+
const bestTexts = ['text4', 'text5', 'text6'];
934+
const [insertResult] = await db.update(usersTable).set({
935+
bestTexts,
936+
}).where(eq(usersTable.id, newUser!.id)).returning();
937+
938+
expect(insertResult?.bestTexts).toEqual(bestTexts);
939+
});
940+
941+
test('insert with array values works', async () => {
942+
const bestTexts = ['text1', 'text2', 'text3'];
943+
const [insertResult] = await db.insert(usersTable).values({
944+
name: 'John',
945+
bestTexts,
946+
}).returning();
947+
948+
expect(insertResult?.bestTexts).toEqual(bestTexts);
949+
});
950+
951+
test('update with array values works', async () => {
952+
const [newUser] = await db.insert(usersTable).values({ name: 'John' }).returning();
953+
954+
const bestTexts = ['text4', 'text5', 'text6'];
955+
const [insertResult] = await db.update(usersTable).set({
956+
bestTexts,
957+
}).where(eq(usersTable.id, newUser!.id)).returning();
958+
959+
expect(insertResult?.bestTexts).toEqual(bestTexts);
960+
});
961+
860962
test('all date and time columns', async () => {
861963
const table = pgTable('all_columns', {
862964
id: serial('id').primaryKey(),

0 commit comments

Comments
 (0)