Skip to content

Commit

Permalink
Merge pull request #1911 from livingforjesus/fix-arrayvalue-issue
Browse files Browse the repository at this point in the history 
[Pg] Add insert/update array support in aws-data-api
  • Loading branch information
@dankochetov
dankochetov committed Apr 10, 2024
2 parents 2c9b73b + 338650d commit 59f2958
Show file tree
Hide file tree
Showing 3 changed files with 171 additions and 14 deletions.
13 changes: 13 additions & 0 deletions drizzle-orm/src/aws-data-api/common/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,19 @@ export function getValueFromDataApi(field: Field) {
if (field.arrayValue.stringValues !== undefined) {
return field.arrayValue.stringValues;
}
if (field.arrayValue.longValues !== undefined) {
return field.arrayValue.longValues;
}
if (field.arrayValue.doubleValues !== undefined) {
return field.arrayValue.doubleValues;
}
if (field.arrayValue.booleanValues !== undefined) {
return field.arrayValue.booleanValues;
}
if (field.arrayValue.arrayValues !== undefined) {
return field.arrayValue.arrayValues;
}

throw new Error('Unknown array type');
} else {
throw new Error('Unknown type');
Expand Down
48 changes: 45 additions & 3 deletions drizzle-orm/src/aws-data-api/pg/driver.ts
View file
Original file line number Diff line number Diff line change
@@ -1,17 +1,20 @@
import { entityKind } from '~/entity.ts';
import type { SQLWrapper } from '~/index.ts';
import { entityKind, is } from '~/entity.ts';
import type { SQL, SQLWrapper } from '~/index.ts';
import { Param, sql, Table } from '~/index.ts';
import type { Logger } from '~/logger.ts';
import { DefaultLogger } from '~/logger.ts';
import { PgDatabase } from '~/pg-core/db.ts';
import { PgDialect } from '~/pg-core/dialect.ts';
import type { PgColumn, PgInsertConfig, PgTable, TableConfig } from '~/pg-core/index.ts';
import { PgArray } from '~/pg-core/index.ts';
import type { PgRaw } from '~/pg-core/query-builders/raw.ts';
import {
createTableRelationsHelpers,
extractTablesRelationalConfig,
type RelationalSchemaConfig,
type TablesRelationalConfig,
} from '~/relations.ts';
import type { DrizzleConfig } from '~/utils.ts';
import type { DrizzleConfig, UpdateSet } from '~/utils.ts';
import type { AwsDataApiClient, AwsDataApiPgQueryResult, AwsDataApiPgQueryResultHKT } from './session.ts';
import { AwsDataApiSession } from './session.ts';

Expand Down Expand Up @@ -48,6 +51,45 @@ export class AwsPgDialect extends PgDialect {
override escapeParam(num: number): string {
return `:${num + 1}`;
}

override buildInsertQuery(
{ table, values, onConflict, returning }: PgInsertConfig<PgTable<TableConfig>>,
): SQL<unknown> {
const columns: Record<string, PgColumn> = table[Table.Symbol.Columns];
const colEntries: [string, PgColumn][] = Object.entries(columns);
for (const value of values) {
for (const [fieldName, col] of colEntries) {
const colValue = value[fieldName];
if (
is(colValue, Param) && colValue.value !== undefined && is(colValue.encoder, PgArray)
&& Array.isArray(colValue.value)
) {
value[fieldName] = sql`cast(${col.mapToDriverValue(colValue.value)} as ${
sql.raw(colValue.encoder.getSQLType())
})`;
}
}
}

return super.buildInsertQuery({ table, values, onConflict, returning });
}

override buildUpdateSet(table: PgTable<TableConfig>, set: UpdateSet): SQL<unknown> {
const columns: Record<string, PgColumn> = table[Table.Symbol.Columns];

for (const [colName, colValue] of Object.entries(set)) {
const currentColumn = columns[colName];
if (
currentColumn && is(colValue, Param) && colValue.value !== undefined && is(colValue.encoder, PgArray)
&& Array.isArray(colValue.value)
) {
set[colName] = sql`cast(${currentColumn?.mapToDriverValue(colValue.value)} as ${
sql.raw(colValue.encoder.getSQLType())
})`;
}
}
return super.buildUpdateSet(table, set);
}
}

export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
Expand Down
124 changes: 113 additions & 11 deletions integration-tests/tests/awsdatapi.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ const usersTable = pgTable('users', {
name: text('name').notNull(),
verified: boolean('verified').notNull().default(false),
jsonb: jsonb('jsonb').$type<string[]>(),
bestTexts: text('best_texts').array().default(sql`'{}'`).notNull(),
createdAt: timestamp('created_at', { withTimezone: true }).notNull().defaultNow(),
});

Expand Down Expand Up @@ -69,6 +70,7 @@ beforeEach(async () => {
name text not null,
verified boolean not null default false,
jsonb jsonb,
best_texts text[] not null default '{}',
created_at timestamptz not null default now()
)
`,
Expand All @@ -84,7 +86,14 @@ test('select all fields', async () => {

expect(result[0]!.createdAt).toBeInstanceOf(Date);
// t.assert(Math.abs(result[0]!.createdAt.getTime() - now) < 100);
expect(result).toEqual([{ id: 1, name: 'John', verified: false, jsonb: null, createdAt: result[0]!.createdAt }]);
expect(result).toEqual([{
bestTexts: [],
id: 1,
name: 'John',
verified: false,
jsonb: null,
createdAt: result[0]!.createdAt,
}]);
});

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

expect(users[0]!.createdAt).toBeInstanceOf(Date);
// t.assert(Math.abs(users[0]!.createdAt.getTime() - now) < 100);
expect(users).toEqual([{ id: 1, name: 'Jane', verified: false, jsonb: null, createdAt: users[0]!.createdAt }]);
expect(users).toEqual([{
id: 1,
bestTexts: [],
name: 'Jane',
verified: false,
jsonb: null,
createdAt: users[0]!.createdAt,
}]);
});

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

expect(users[0]!.createdAt).toBeInstanceOf(Date);
// t.assert(Math.abs(users[0]!.createdAt.getTime() - now) < 100);
expect(users).toEqual([{ id: 1, name: 'John', verified: false, jsonb: null, createdAt: users[0]!.createdAt }]);
expect(users).toEqual([{
bestTexts: [],
id: 1,
name: 'John',
verified: false,
jsonb: null,
createdAt: users[0]!.createdAt,
}]);
});

test('delete with returning partial', async () => {
Expand All @@ -211,13 +234,20 @@ test('delete with returning partial', async () => {
test('insert + select', async () => {
await db.insert(usersTable).values({ name: 'John' });
const result = await db.select().from(usersTable);
expect(result).toEqual([{ id: 1, name: 'John', verified: false, jsonb: null, createdAt: result[0]!.createdAt }]);
expect(result).toEqual([{
bestTexts: [],
id: 1,
name: 'John',
verified: false,
jsonb: null,
createdAt: result[0]!.createdAt,
}]);

await db.insert(usersTable).values({ name: 'Jane' });
const result2 = await db.select().from(usersTable);
expect(result2).toEqual([
{ id: 1, name: 'John', verified: false, jsonb: null, createdAt: result2[0]!.createdAt },
{ id: 2, name: 'Jane', verified: false, jsonb: null, createdAt: result2[1]!.createdAt },
{ bestTexts: [], id: 1, name: 'John', verified: false, jsonb: null, createdAt: result2[0]!.createdAt },
{ bestTexts: [], id: 2, name: 'Jane', verified: false, jsonb: null, createdAt: result2[1]!.createdAt },
]);
});

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

expect(result).toEqual([{ id: 1, name: 'John', verified: true, jsonb: null, createdAt: result[0]!.createdAt }]);
expect(result).toEqual([{
bestTexts: [],
id: 1,
name: 'John',
verified: true,
jsonb: null,
createdAt: result[0]!.createdAt,
}]);
});

test('insert many', async () => {
Expand Down Expand Up @@ -385,12 +422,14 @@ test('full join with alias', async () => {
expect(result).toEqual([{
users: {
id: 10,
bestTexts: [],
name: 'Ivan',
verified: false,
jsonb: null,
createdAt: result[0]!.users.createdAt,
},
customer: {
bestTexts: [],
id: 11,
name: 'Hans',
verified: false,
Expand Down Expand Up @@ -627,7 +666,7 @@ test('build query insert with onConflict do update', async () => {

expect(query).toEqual({
sql:
'insert into "users" ("id", "name", "verified", "jsonb", "created_at") values (default, :1, default, :2, default) on conflict ("id") do update set "name" = :3',
'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',
params: ['John', '["foo","bar"]', 'John1'],
// typings: ['none', 'json', 'none']
});
Expand All @@ -641,7 +680,7 @@ test('build query insert with onConflict do update / multiple columns', async ()

expect(query).toEqual({
sql:
'insert into "users" ("id", "name", "verified", "jsonb", "created_at") values (default, :1, default, :2, default) on conflict ("id","name") do update set "name" = :3',
'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',
params: ['John', '["foo","bar"]', 'John1'],
// typings: ['none', 'json', 'none']
});
Expand All @@ -655,7 +694,7 @@ test('build query insert with onConflict do nothing', async () => {

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

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

test('insert with array values works', async () => {
const bestTexts = ['text1', 'text2', 'text3'];
const [insertResult] = await db.insert(usersTable).values({
name: 'John',
bestTexts,
}).returning();

expect(insertResult?.bestTexts).toEqual(bestTexts);
});

test('update with array values works', async () => {
const [newUser] = await db.insert(usersTable).values({ name: 'John' }).returning();

const bestTexts = ['text4', 'text5', 'text6'];
const [insertResult] = await db.update(usersTable).set({
bestTexts,
}).where(eq(usersTable.id, newUser!.id)).returning();

expect(insertResult?.bestTexts).toEqual(bestTexts);
});

test('insert with array values works', async () => {
const bestTexts = ['text1', 'text2', 'text3'];
const [insertResult] = await db.insert(usersTable).values({
name: 'John',
bestTexts,
}).returning();

expect(insertResult?.bestTexts).toEqual(bestTexts);
});

test('update with array values works', async () => {
const [newUser] = await db.insert(usersTable).values({ name: 'John' }).returning();

const bestTexts = ['text4', 'text5', 'text6'];
const [insertResult] = await db.update(usersTable).set({
bestTexts,
}).where(eq(usersTable.id, newUser!.id)).returning();

expect(insertResult?.bestTexts).toEqual(bestTexts);
});

test('insert with array values works', async () => {
const bestTexts = ['text1', 'text2', 'text3'];
const [insertResult] = await db.insert(usersTable).values({
name: 'John',
bestTexts,
}).returning();

expect(insertResult?.bestTexts).toEqual(bestTexts);
});

test('update with array values works', async () => {
const [newUser] = await db.insert(usersTable).values({ name: 'John' }).returning();

const bestTexts = ['text4', 'text5', 'text6'];
const [insertResult] = await db.update(usersTable).set({
bestTexts,
}).where(eq(usersTable.id, newUser!.id)).returning();

expect(insertResult?.bestTexts).toEqual(bestTexts);
});

test('all date and time columns', async () => {
const table = pgTable('all_columns', {
id: serial('id').primaryKey(),
Expand Down

0 comments on commit 59f2958

Please sign in to comment.