[Pg] Add insert/update array support in aws-data-api

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

@@ -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');

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

@@ -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';
 
@@ -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>>(

integration-tests/tests/awsdatapi.test.ts

@@ -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(),
 });
 
@@ -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()
 			)
 		`,
@@ -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 () => {
@@ -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 () => {
@@ -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 () => {
@@ -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 },
 	]);
 });
 
@@ -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 () => {
@@ -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,
@@ -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']
 	});
@@ -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']
 	});
@@ -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']
 	});
@@ -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']
 	});
@@ -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(),