Merge pull request #2241 from drizzle-team/extend-proxy

Extend proxy

changelogs/drizzle-orm/0.30.10.md

@@ -0,0 +1,17 @@
+## New Features
+
+### 🎉 `.if()` function added to all WHERE expressions
+
+#### Select all users after cursors if a cursor value was provided
+
+```ts
+function getUsersAfter(cursor?: number) {
+    return db.select().from(users).where(
+        gt(users.id, cursor).if(cursor)
+    );
+}
+```
+
+## Bug Fixes
+
+- Fixed internal mappings for sessions `.all`, `.values`, `.execute` functions in AWS DataAPI

drizzle-orm/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "drizzle-orm",
-	"version": "0.30.9",
+	"version": "0.30.10",
 	"description": "Drizzle ORM package for SQL databases",
 	"type": "module",
 	"scripts": {

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

@@ -24,7 +24,9 @@ import { getValueFromDataApi, toValueParam } from '../common/index.ts';
 
 export type AwsDataApiClient = RDSDataClient;
 
-export class AwsDataApiPreparedQuery<T extends PreparedQueryConfig> extends PgPreparedQuery<T> {
+export class AwsDataApiPreparedQuery<
+	T extends PreparedQueryConfig & { values: AwsDataApiPgQueryResult<unknown[]> },
+> extends PgPreparedQuery<T> {
 	static readonly [entityKind]: string = 'AwsDataApiPreparedQuery';
 
 	private rawQuery: ExecuteStatementCommand;
@@ -56,18 +58,44 @@ export class AwsDataApiPreparedQuery<T extends PreparedQueryConfig> extends PgPr
 	async execute(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['execute']> {
 		const { fields, joinsNotNullableMap, customResultMapper } = this;
 
-		const result = await this.values(placeholderValues) as AwsDataApiPgQueryResult<unknown[]>;
+		const result = await this.values(placeholderValues);
 		if (!fields && !customResultMapper) {
-			return result as T['execute'];
+			const { columnMetadata, rows } = result;
+			if (!columnMetadata) {
+				return result;
+			}
+			const mappedRows = rows.map((sourceRow) => {
+				const row: Record<string, unknown> = {};
+				for (const [index, value] of sourceRow.entries()) {
+					const metadata = columnMetadata[index];
+					if (!metadata) {
+						throw new Error(
+							`Unexpected state: no column metadata found for index ${index}. Please report this issue on GitHub: https://github.com/drizzle-team/drizzle-orm/issues/new/choose`,
+						);
+					}
+					if (!metadata.name) {
+						throw new Error(
+							`Unexpected state: no column name for index ${index} found in the column metadata. Please report this issue on GitHub: https://github.com/drizzle-team/drizzle-orm/issues/new/choose`,
+						);
+					}
+					row[metadata.name] = value;
+				}
+				return row;
+			});
+			return Object.assign(result, { rows: mappedRows });
 		}
+
 		return customResultMapper
 			? customResultMapper(result.rows!)
-			: result.rows!.map((row) => mapResultRow<T['execute']>(fields!, row, joinsNotNullableMap));
+			: result.rows!.map((row) => mapResultRow(fields!, row, joinsNotNullableMap));
 	}
 
 	async all(placeholderValues?: Record<string, unknown> | undefined): Promise<T['all']> {
-		const result = await this.execute(placeholderValues) as AwsDataApiPgQueryResult<unknown>;
-		return result.rows;
+		const result = await this.execute(placeholderValues);
+		if (!this.fields && !this.customResultMapper) {
+			return (result as AwsDataApiPgQueryResult<unknown>).rows;
+		}
+		return result;
 	}
 
 	async values(placeholderValues: Record<string, unknown> = {}): Promise<T['values']> {
@@ -80,20 +108,7 @@ export class AwsDataApiPreparedQuery<T extends PreparedQueryConfig> extends PgPr
 
 		this.options.logger?.logQuery(this.rawQuery.input.sql!, this.rawQuery.input.parameters);
 
-		const { fields, rawQuery, client, customResultMapper } = this;
-		if (!fields && !customResultMapper) {
-			const result = await client.send(rawQuery);
-			if (result.columnMetadata && result.columnMetadata.length > 0) {
-				const rows = this.mapResultRows(result.records ?? [], result.columnMetadata);
-				return {
-					...result,
-					rows,
-				};
-			}
-			return result;
-		}
-
-		const result = await client.send(rawQuery);
+		const result = await this.client.send(this.rawQuery);
 		const rows = result.records?.map((row) => {
 			return row.map((field) => getValueFromDataApi(field));
 		}) ?? [];
@@ -161,14 +176,20 @@ export class AwsDataApiSession<
 		};
 	}
 
-	prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
+	prepareQuery<
+		T extends PreparedQueryConfig & {
+			values: AwsDataApiPgQueryResult<unknown[]>;
+		} = PreparedQueryConfig & {
+			values: AwsDataApiPgQueryResult<unknown[]>;
+		},
+	>(
 		query: QueryWithTypings,
 		fields: SelectedFieldsOrdered | undefined,
 		name: string | undefined,
 		isResponseInArrayMode: boolean,
 		customResultMapper?: (rows: unknown[][]) => T['execute'],
 		transactionId?: string,
-	): PgPreparedQuery<T> {
+	): AwsDataApiPreparedQuery<T> {
 		return new AwsDataApiPreparedQuery(
 			this.client,
 			query.sql,
@@ -183,7 +204,7 @@ export class AwsDataApiSession<
 	}
 
 	override execute<T>(query: SQL): Promise<T> {
-		return this.prepareQuery<PreparedQueryConfig & { execute: T }>(
+		return this.prepareQuery<PreparedQueryConfig & { execute: T; values: AwsDataApiPgQueryResult<unknown[]> }>(
 			this.dialect.sqlToQuery(query),
 			undefined,
 			undefined,

drizzle-orm/src/pg-proxy/driver.ts

@@ -18,13 +18,15 @@ export type RemoteCallback = (
 	sql: string,
 	params: any[],
 	method: 'all' | 'execute',
+	typings?: any[],
 ) => Promise<{ rows: any[] }>;
 
 export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
 	callback: RemoteCallback,
 	config: DrizzleConfig<TSchema> = {},
+	_dialect: () => PgDialect = () => new PgDialect(),
 ): PgRemoteDatabase<TSchema> {
-	const dialect = new PgDialect();
+	const dialect = _dialect();
 	let logger;
 	if (config.logger === true) {
 		logger = new DefaultLogger();

drizzle-orm/src/pg-proxy/session.ts

@@ -7,7 +7,8 @@ import type { SelectedFieldsOrdered } from '~/pg-core/query-builders/select.type
 import type { PgTransactionConfig, PreparedQueryConfig, QueryResultHKT } from '~/pg-core/session.ts';
 import { PgPreparedQuery as PreparedQueryBase, PgSession } from '~/pg-core/session.ts';
 import type { RelationalSchemaConfig, TablesRelationalConfig } from '~/relations.ts';
-import { fillPlaceholders, type Query } from '~/sql/sql.ts';
+import type { QueryWithTypings } from '~/sql/sql.ts';
+import { fillPlaceholders } from '~/sql/sql.ts';
 import { tracer } from '~/tracing.ts';
 import { type Assume, mapResultRow } from '~/utils.ts';
 import type { RemoteCallback } from './driver.ts';
@@ -35,7 +36,7 @@ export class PgRemoteSession<
 	}
 
 	prepareQuery<T extends PreparedQueryConfig>(
-		query: Query,
+		query: QueryWithTypings,
 		fields: SelectedFieldsOrdered | undefined,
 		name: string | undefined,
 		isResponseInArrayMode: boolean,
@@ -45,6 +46,7 @@ export class PgRemoteSession<
 			this.client,
 			query.sql,
 			query.params,
+			query.typings,
 			this.logger,
 			fields,
 			isResponseInArrayMode,
@@ -80,6 +82,7 @@ export class PreparedQuery<T extends PreparedQueryConfig> extends PreparedQueryB
 		private client: RemoteCallback,
 		private queryString: string,
 		private params: unknown[],
+		private typings: any[] | undefined,
 		private logger: Logger,
 		private fields: SelectedFieldsOrdered | undefined,
 		private _isResponseInArrayMode: boolean,
@@ -91,7 +94,7 @@ export class PreparedQuery<T extends PreparedQueryConfig> extends PreparedQueryB
 	async execute(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['execute']> {
 		return tracer.startActiveSpan('drizzle.execute', async (span) => {
 			const params = fillPlaceholders(this.params, placeholderValues);
-			const { fields, client, queryString, joinsNotNullableMap, customResultMapper, logger } = this;
+			const { fields, client, queryString, joinsNotNullableMap, customResultMapper, logger, typings } = this;
 
 			span?.setAttributes({
 				'drizzle.query.text': queryString,
@@ -102,7 +105,7 @@ export class PreparedQuery<T extends PreparedQueryConfig> extends PreparedQueryB
 
 			if (!fields && !customResultMapper) {
 				return tracer.startActiveSpan('drizzle.driver.execute', async () => {
-					const { rows } = await client(queryString, params as any[], 'execute');
+					const { rows } = await client(queryString, params as any[], 'execute', typings);
 
 					return rows;
 				});
@@ -114,7 +117,7 @@ export class PreparedQuery<T extends PreparedQueryConfig> extends PreparedQueryB
 					'drizzle.query.params': JSON.stringify(params),
 				});
 
-				const { rows } = await client(queryString, params as any[], 'all');
+				const { rows } = await client(queryString, params as any[], 'all', typings);
 
 				return rows;
 			});