Improve batch types

drizzle-orm/src/d1/driver.ts

@@ -22,7 +22,7 @@ export class DrizzleD1Database<
 	/** @internal */
 	declare readonly session: SQLiteD1Session<TSchema, ExtractTablesWithRelations<TSchema>>;
 
-	async batch<U extends BatchItem<'sqlite'>, T extends Readonly<[U, ...U[]]>>(
+	async batch<U extends BatchItem<'sqlite'>, T extends U[] | Readonly<[U, ...U[]]>>(
 		batch: T,
 	): Promise<BatchResponse<T>> {
 		return this.session.batch(batch) as Promise<BatchResponse<T>>;

drizzle-orm/src/d1/session.ts

@@ -61,7 +61,7 @@ export class SQLiteD1Session<
 		);
 	}
 
-	async batch<U extends BatchItem, T extends Readonly<[U, ...U[]]>>(queries: T) {
+	async batch<U extends BatchItem<'sqlite'>, T extends U[] | Readonly<[U, ...U[]]>>(queries: T) {
 		const preparedQueries: PreparedQuery[] = [];
 		const builtQueries: D1PreparedStatement[] = [];
 

drizzle-orm/src/libsql/driver.ts

@@ -22,7 +22,7 @@ export class LibSQLDatabase<
 	/** @internal */
 	declare readonly session: LibSQLSession<TSchema, ExtractTablesWithRelations<TSchema>>;
 
-	async batch<U extends BatchItem<'sqlite'>, T extends Readonly<[U, ...U[]]>>(
+	async batch<U extends BatchItem<'sqlite'>, T extends U[] | Readonly<[U, ...U[]]>>(
 		batch: T,
 	): Promise<BatchResponse<T>> {
 		return this.session.batch(batch) as Promise<BatchResponse<T>>;

drizzle-orm/src/libsql/session.ts

@@ -61,7 +61,7 @@ export class LibSQLSession<
 		);
 	}
 
-	async batch<T extends BatchItem<'sqlite'>[] | readonly BatchItem<'sqlite'>[]>(queries: T) {
+	async batch<U extends BatchItem<'sqlite'>, T extends U[] | Readonly<[U, ...U[]]>>(queries: T) {
 		const preparedQueries: PreparedQuery[] = [];
 		const builtQueries: InStatement[] = [];
 

drizzle-orm/src/neon-http/driver.ts

@@ -47,7 +47,7 @@ export class NeonHttpDatabase<
 	/** @internal */
 	declare readonly session: NeonHttpSession<TSchema, ExtractTablesWithRelations<TSchema>>;
 
-	async batch<U extends BatchItem<'pg'>, T extends Readonly<[U, ...U[]]>>(
+	async batch<U extends BatchItem<'pg'>, T extends U[] | Readonly<[U, ...U[]]>>(
 		batch: T,
 	): Promise<BatchResponse<T>> {
 		return this.session.batch(batch) as Promise<BatchResponse<T>>;

drizzle-orm/src/neon-http/session.ts

@@ -125,7 +125,7 @@ export class NeonHttpSession<
 		);
 	}
 
-	async batch<U extends BatchItem<'pg'>, T extends Readonly<[U, ...U[]]>>(queries: T) {
+	async batch<U extends BatchItem<'pg'>, T extends U[] | Readonly<[U, ...U[]]>>(queries: T) {
 		const preparedQueries: PreparedQuery[] = [];
 		const builtQueries: NeonQueryPromise<any, true>[] = [];
 

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

@@ -20,7 +20,7 @@ export class SqliteRemoteDatabase<
 	/** @internal */
 	declare readonly session: SQLiteRemoteSession<TSchema, ExtractTablesWithRelations<TSchema>>;
 
-	async batch<U extends BatchItem<'sqlite'>, T extends Readonly<[U, ...U[]]>>(
+	async batch<U extends BatchItem<'sqlite'>, T extends U[] | Readonly<[U, ...U[]]>>(
 		batch: T,
 	): Promise<BatchResponse<T>> {
 		return this.session.batch(batch) as Promise<BatchResponse<T>>;

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

@@ -35,7 +35,7 @@ export class SQLiteRemoteSession<
 		private client: RemoteCallback,
 		dialect: SQLiteAsyncDialect,
 		private schema: RelationalSchemaConfig<TSchema> | undefined,
-		private batchCLient?: AsyncBatchRemoteCallback,
+		private batchClient?: AsyncBatchRemoteCallback,
 		options: SQLiteRemoteSessionOptions = {},
 	) {
 		super(dialect);
@@ -60,7 +60,7 @@ export class SQLiteRemoteSession<
 		);
 	}
 
-	async batch<T extends BatchItem<'sqlite'>[] | readonly BatchItem<'sqlite'>[]>(queries: T) {
+	async batch<U extends BatchItem<'sqlite'>, T extends U[] | Readonly<[U, ...U[]]>>(queries: T) {
 		const preparedQueries: PreparedQuery[] = [];
 		const builtQueries: { sql: string; params: any[]; method: 'run' | 'all' | 'values' | 'get' }[] = [];
 
@@ -71,7 +71,7 @@ export class SQLiteRemoteSession<
 			builtQueries.push({ sql: builtQuery.sql, params: builtQuery.params, method: builtQuery.method });
 		}
 
-		const batchResults = await (this.batchCLient as AsyncBatchRemoteCallback)(builtQueries);
+		const batchResults = await (this.batchClient as AsyncBatchRemoteCallback)(builtQueries);
 		return batchResults.map((result, i) => preparedQueries[i]!.mapResult(result, true));
 	}
 

integration-tests/tests/d1-batch.test.ts

@@ -545,6 +545,59 @@ test('insert + delete + select + select partial', async () => {
 	);
 });
 
+test('dynamic array select', async () => {
+	await db.insert(usersTable).values([
+		{ id: 1, name: 'John' },
+		{ id: 2, name: 'Dan' },
+	]);
+
+	const batchResponse = await db.batch(
+		[1, 2].map((id) => db.query.usersTable.findFirst({
+			where: eq(usersTable.id, id),
+		})),
+	);
+
+	expectTypeOf(batchResponse).toEqualTypeOf<({
+		id: number;
+		name: string;
+		verified: number;
+		invitedBy: number | null;
+	} | undefined)[]>();
+
+	expect(batchResponse.length).eq(2);
+
+	expect(batchResponse[0]).toEqual([{ id: 1 }]);
+
+	expect(batchResponse[1]).toEqual([{ id: 2 }]);
+});
+
+test('dynamic array insert + select', async () => {
+	const queries = [];
+
+	queries.push(db.insert(usersTable).values([
+		{ id: 1, name: 'John' },
+		{ id: 2, name: 'Dan' },
+	]));
+	queries.push(db.query.usersTable.findMany({}));
+
+	const batchResponse = await db.batch(queries);
+
+	expectTypeOf(batchResponse).toEqualTypeOf<(
+		D1Result | {
+			id: number;
+			name: string;
+			verified: number;
+			invitedBy: number | null;
+		}[]
+	)[]>();
+
+	expect(batchResponse.length).eq(2);
+
+	expect(batchResponse[0]).toEqual({ changes: 2, lastInsertRowid: 2 });
+
+	expect(batchResponse[1]).toEqual([{ id: 1 }, { id: 2 }]);
+});
+
 // * additionally
 // batch for all libsql cases, just replace simple calls with batch calls
 // batch for all rqb cases, just replace simple calls with batch calls

integration-tests/tests/libsql-batch.test.ts

@@ -564,6 +564,59 @@ test('insert + delete + select + select partial', async () => {
 	);
 });
 
+test('dynamic array select', async () => {
+	await db.insert(usersTable).values([
+		{ id: 1, name: 'John' },
+		{ id: 2, name: 'Dan' },
+	]);
+
+	const batchResponse = await db.batch(
+		[1, 2].map((id) => db.query.usersTable.findFirst({
+			where: eq(usersTable.id, id),
+		})),
+	);
+
+	expectTypeOf(batchResponse).toEqualTypeOf<({
+		id: number;
+		name: string;
+		verified: number;
+		invitedBy: number | null;
+	} | undefined)[]>();
+
+	expect(batchResponse.length).eq(2);
+
+	expect(batchResponse[0]).toEqual([{ id: 1 }]);
+
+	expect(batchResponse[1]).toEqual([{ id: 2 }]);
+});
+
+test('dynamic array insert + select', async () => {
+	const queries = [];
+
+	queries.push(db.insert(usersTable).values([
+		{ id: 1, name: 'John' },
+		{ id: 2, name: 'Dan' },
+	]));
+	queries.push(db.query.usersTable.findMany({}));
+
+	const batchResponse = await db.batch(queries);
+
+	expectTypeOf(batchResponse).toEqualTypeOf<(
+		ResultSet | {
+			id: number;
+			name: string;
+			verified: number;
+			invitedBy: number | null;
+		}[]
+	)[]>();
+
+	expect(batchResponse.length).eq(2);
+
+	expect(batchResponse[0]).toEqual({ changes: 2, lastInsertRowid: 2 });
+
+	expect(batchResponse[1]).toEqual([{ id: 1 }, { id: 2 }]);
+});
+
 // * additionally
 // batch for all libsql cases, just replace simple calls with batch calls
 // batch for all rqb cases, just replace simple calls with batch calls

integration-tests/tests/neon-http-batch.test.ts

@@ -551,6 +551,59 @@ test('select raw', async () => {
 	});
 });
 
+test('dynamic array select', async () => {
+	await db.insert(usersTable).values([
+		{ id: 1, name: 'John' },
+		{ id: 2, name: 'Dan' },
+	]);
+
+	const batchResponse = await db.batch(
+		[1, 2].map((id) => db.query.usersTable.findFirst({
+			where: eq(usersTable.id, id),
+		})),
+	);
+
+	expectTypeOf(batchResponse).toEqualTypeOf<({
+		id: number;
+		name: string;
+		verified: number;
+		invitedBy: number | null;
+	} | undefined)[]>();
+
+	expect(batchResponse.length).eq(2);
+
+	expect(batchResponse[0]).toEqual([{ id: 1 }]);
+
+	expect(batchResponse[1]).toEqual([{ id: 2 }]);
+});
+
+test('dynamic array insert + select', async () => {
+	const queries = [];
+
+	queries.push(db.insert(usersTable).values([
+		{ id: 1, name: 'John' },
+		{ id: 2, name: 'Dan' },
+	]));
+	queries.push(db.query.usersTable.findMany({}));
+
+	const batchResponse = await db.batch(queries);
+
+	expectTypeOf(batchResponse).toEqualTypeOf<(
+		NeonHttpQueryResult<never> | {
+			id: number;
+			name: string;
+			verified: number;
+			invitedBy: number | null;
+		}[]
+	)[]>();
+
+	expect(batchResponse.length).eq(2);
+
+	expect(batchResponse[0]).toEqual({ changes: 2, lastInsertRowid: 2 });
+
+	expect(batchResponse[1]).toEqual([{ id: 1 }, { id: 2 }]);
+});
+
 // * additionally
 // batch for all neon cases, just replace simple calls with batch calls
 // batch for all rqb cases, just replace simple calls with batch calls

integration-tests/tests/sqlite-proxy-batch.test.ts

@@ -700,3 +700,56 @@ test('insert + delete + select + select partial', async () => {
 		{ id: 2, invitedBy: null },
 	);
 });
+
+test('dynamic array select', async () => {
+	await db.insert(usersTable).values([
+		{ id: 1, name: 'John' },
+		{ id: 2, name: 'Dan' },
+	]);
+
+	const batchResponse = await db.batch(
+		[1, 2].map((id) => db.query.usersTable.findFirst({
+			where: eq(usersTable.id, id),
+		})),
+	);
+
+	expectTypeOf(batchResponse).toEqualTypeOf<({
+		id: number;
+		name: string;
+		verified: number;
+		invitedBy: number | null;
+	} | undefined)[]>();
+
+	expect(batchResponse.length).eq(2);
+
+	expect(batchResponse[0]).toEqual([{ id: 1 }]);
+
+	expect(batchResponse[1]).toEqual([{ id: 2 }]);
+});
+
+test('dynamic array insert + select', async () => {
+	const queries = [];
+
+	queries.push(db.insert(usersTable).values([
+		{ id: 1, name: 'John' },
+		{ id: 2, name: 'Dan' },
+	]));
+	queries.push(db.query.usersTable.findMany({}));
+
+	const batchResponse = await db.batch(queries);
+
+	expectTypeOf(batchResponse).toEqualTypeOf<(
+		SqliteRemoteResult | {
+			id: number;
+			name: string;
+			verified: number;
+			invitedBy: number | null;
+		}[]
+	)[]>();
+
+	expect(batchResponse.length).eq(2);
+
+	expect(batchResponse[0]).toEqual({ changes: 2, lastInsertRowid: 2 });
+
+	expect(batchResponse[1]).toEqual([{ id: 1 }, { id: 2 }]);
+});