Fixed missing bun drivers from kit, fixed lack of subquery procession in cockroach, mssql sessions, fixed broken test cases

Author: Sukairo-02

drizzle-kit/src/cli/connections.ts

@@ -44,7 +44,8 @@ export const preparePostgresDB = async (
 			| 'pg'
 			| 'postgres'
 			| '@vercel/postgres'
-			| '@neondatabase/serverless';
+			| '@neondatabase/serverless'
+			| 'bun';
 		proxy: Proxy;
 		transactionProxy: TransactionProxy;
 		migrate: (config: string | MigrationConfig) => Promise<void | MigratorInitFailResponse>;
@@ -559,8 +560,64 @@ export const preparePostgresDB = async (
 		return { packageName: '@neondatabase/serverless', query, proxy, transactionProxy, migrate: migrateFn };
 	}
 
+	if (await checkPackage('bun')) {
+		console.log(withStyle.info(`Using 'bun' driver for database querying`));
+		const { SQL } = await import('bun');
+		const { drizzle } = await import('drizzle-orm/bun-sql');
+		const { migrate } = await import('drizzle-orm/bun-sql/migrator');
+
+		const ssl = 'ssl' in credentials
+			? credentials.ssl === 'prefer'
+					|| credentials.ssl === 'require'
+					|| credentials.ssl === 'allow'
+				? true
+				: false
+			: undefined;
+
+		const client = new SQL({
+			adapter: 'postgres',
+			...credentials,
+			ssl,
+			max: 1,
+		});
+		const db = drizzle({ client });
+		const migrateFn = async (config: MigrationConfig) => {
+			return migrate(db, config);
+		};
+
+		const query = async (sql: string, params?: any[]) => {
+			const result = await client.unsafe(sql, params ?? []);
+			return result;
+		};
+
+		const proxy: Proxy = async (params) => {
+			const query = client.unsafe(params.sql, params.params);
+			if (params.mode === 'array') {
+				return await query.values();
+			}
+			return await query;
+		};
+
+		const transactionProxy: TransactionProxy = async (queries) => {
+			const results: any[] = [];
+			try {
+				await client.transaction(async (tx) => {
+					for (const query of queries) {
+						const result = await tx.unsafe(query.sql);
+						results.push(result);
+					}
+				});
+			} catch (error) {
+				results.push(error as Error);
+			}
+			return results;
+		};
+
+		return { packageName: 'bun', query, proxy, transactionProxy, migrate: migrateFn };
+	}
+
 	console.error(
-		"To connect to Postgres database - please install either of 'pg', 'postgres', '@neondatabase/serverless' or '@vercel/postgres' drivers",
+		"To connect to Postgres database - please install either of 'pg', 'postgres', 'bun', '@neondatabase/serverless' or '@vercel/postgres' drivers",
 	);
 	process.exit(1);
 };
@@ -859,7 +916,7 @@ export const connectToMySQL = async (
 	it: MysqlCredentials,
 ): Promise<{
 	db: DB;
-	packageName: 'mysql2' | '@planetscale/database';
+	packageName: 'mysql2' | '@planetscale/database' | 'bun';
 	proxy: Proxy;
 	transactionProxy: TransactionProxy;
 	database: string;
@@ -868,6 +925,7 @@ export const connectToMySQL = async (
 	const result = parseMysqlCredentials(it);
 
 	if (await checkPackage('mysql2')) {
+		console.log(withStyle.info(`Using 'mysql2' driver for database querying`));
 		const { createConnection } = await import('mysql2/promise');
 		const { drizzle } = await import('drizzle-orm/mysql2');
 		const { migrate } = await import('drizzle-orm/mysql2/migrator');
@@ -945,6 +1003,7 @@ export const connectToMySQL = async (
 	}
 
 	if (await checkPackage('@planetscale/database')) {
+		console.log(withStyle.info(`Using '@planetscale/database' driver for database querying`));
 		const { Client } = await import('@planetscale/database');
 		const { drizzle } = await import('drizzle-orm/planetscale-serverless');
 		const { migrate } = await import(
@@ -1000,8 +1059,73 @@ export const connectToMySQL = async (
 		};
 	}
 
+	if (await checkPackage('bun')) {
+		console.log(withStyle.info(`Using 'bun' driver for database querying`));
+		const { SQL } = await import('bun');
+		const { drizzle } = await import('drizzle-orm/bun-sql');
+		const { migrate } = await import('drizzle-orm/bun-sql/migrator');
+
+		const ssl = result.credentials && 'ssl' in result.credentials
+			? result.credentials.ssl === 'prefer'
+					|| result.credentials.ssl === 'require'
+					|| result.credentials.ssl === 'allow'
+				? true
+				: false
+			: undefined;
+
+		const client = result.url
+			? new SQL(result.url)
+			: new SQL({
+				adapter: 'mysql',
+				...result.credentials,
+				ssl,
+			});
+
+		const db = drizzle({ client });
+		const migrateFn = async (config: MigrationConfig) => {
+			return migrate(db, config);
+		};
+
+		const query = async (sql: string, params?: any[]) => {
+			const result = await client.unsafe(sql, params ?? []);
+			return result;
+		};
+
+		const proxy: Proxy = async (params) => {
+			const query = client.unsafe(params.sql, params.params);
+			if (params.mode === 'array') {
+				return await query.values();
+			}
+			return await query;
+		};
+
+		const transactionProxy: TransactionProxy = async (queries) => {
+			const results: any[] = [];
+			try {
+				await client.transaction(async (tx) => {
+					for (const query of queries) {
+						const result = await tx.unsafe(query.sql);
+						results.push(result);
+					}
+				});
+			} catch (error) {
+				results.push(error as Error);
+			}
+			return results;
+		};
+
+		return {
+			packageName: 'bun',
+			db: { query },
+			proxy,
+			transactionProxy,
+			migrate: migrateFn,
+			database: result.database,
+		};
+	}
+
 	console.error(
-		"To connect to MySQL database - please install either of 'mysql2' or '@planetscale/database' drivers",
+		"To connect to MySQL database - please install either of 'mysql2', 'bun' or '@planetscale/database' drivers",
 	);
 	process.exit(1);
 };
@@ -1129,7 +1253,13 @@ export const connectToSQLite = async (
 ): Promise<
 	& SQLiteDB
 	& {
-		packageName: 'd1-http' | '@libsql/client' | 'better-sqlite3' | '@sqlitecloud/drivers' | '@tursodatabase/database';
+		packageName:
+			| 'd1-http'
+			| '@libsql/client'
+			| 'better-sqlite3'
+			| '@sqlitecloud/drivers'
+			| '@tursodatabase/database'
+			| 'bun';
 		migrate: (config: string | MigrationConfig) => Promise<void | MigratorInitFailResponse>;
 		proxy: Proxy;
 		transactionProxy: TransactionProxy;
@@ -1246,15 +1376,14 @@ export const connectToSQLite = async (
 				);
 			};
 
-			const db: SQLiteDB = {
-				query: async <T>(sql: string, params?: any[]) => {
-					const res = await remoteCallback(sql, params || [], 'all');
-					return res.rows as T[];
-				},
-				run: async (query: string) => {
-					await remoteCallback(query, [], 'run');
-				},
+			const query = async <T>(sql: string, params?: any[]) => {
+				const res = await remoteCallback(sql, params || [], 'all');
+				return res.rows as T[];
+			};
+			const run = async (query: string) => {
+				await remoteCallback(query, [], 'run');
 			};
+
 			const proxy: Proxy = async (params) => {
 				const preparedParams = prepareSqliteParams(params.params || [], 'd1-http');
 				const result = await remoteCallback(
@@ -1269,7 +1398,7 @@ export const connectToSQLite = async (
 				const result = await remoteBatchCallback(queries);
 				return result.rows;
 			};
-			return { ...db, packageName: 'd1-http', proxy, transactionProxy, migrate: migrateFn };
+			return { query, run, packageName: 'd1-http', proxy, transactionProxy, migrate: migrateFn };
 		} else if (driver === 'sqlite-cloud') {
 			assertPackages('@sqlitecloud/drivers');
 			const { Database } = await import('@sqlitecloud/drivers');
@@ -1360,6 +1489,7 @@ export const connectToSQLite = async (
 	}
 
 	if (await checkPackage('@libsql/client')) {
+		console.log(withStyle.info(`Using '@libsql/client' driver for database querying`));
 		const { createClient } = await import('@libsql/client');
 		const { drizzle } = await import('drizzle-orm/libsql');
 		const { migrate } = await import('drizzle-orm/libsql/migrator');
@@ -1372,18 +1502,12 @@ export const connectToSQLite = async (
 			return migrate(drzl, config);
 		};
 
-		const db: SQLiteDB = {
-			query: async <T>(sql: string, params?: any[]) => {
-				const res = await client.execute({ sql, args: params || [] }).catch((e) => {
-					throw new QueryError(e, sql, params || []);
-				});
-				return res.rows as T[];
-			},
-			run: async (query: string) => {
-				await client.execute(query).catch((e) => {
-					throw new QueryError(e, query, []);
-				});
-			},
+		const query = async <T>(sql: string, params?: any[]) => {
+			const res = await client.execute({ sql, args: params || [] });
+			return res.rows as T[];
+		};
+		const run = async (query: string) => {
+			await client.execute(query);
 		};
 
 		type Transaction = Awaited<ReturnType<typeof client.transaction>>;
@@ -1423,7 +1547,7 @@ export const connectToSQLite = async (
 			return results;
 		};
 
-		return { ...db, packageName: '@libsql/client', proxy, transactionProxy, migrate: migrateFn };
+		return { query, run, packageName: '@libsql/client', proxy, transactionProxy, migrate: migrateFn };
 	}
 
 	if (await checkPackage('@tursodatabase/database')) {
@@ -1482,6 +1606,7 @@ export const connectToSQLite = async (
 	}
 
 	if (await checkPackage('better-sqlite3')) {
+		console.log(withStyle.info(`Using 'better-sqlite3' driver for database querying`));
 		const { default: Database } = await import('better-sqlite3');
 		const { drizzle } = await import('drizzle-orm/better-sqlite3');
 		const { migrate } = await import('drizzle-orm/better-sqlite3/migrator');
@@ -1550,8 +1675,65 @@ export const connectToSQLite = async (
 		return { ...db, packageName: 'better-sqlite3', proxy, transactionProxy, migrate: migrateFn };
 	}
 
+	if (await checkPackage('bun')) {
+		console.log(withStyle.info(`Using 'bun' driver for database querying`));
+		const { SQL } = await import('bun');
+		const { drizzle } = await import('drizzle-orm/bun-sql');
+		const { migrate } = await import('drizzle-orm/bun-sql/migrator');
+
+		const client = new SQL({
+			adapter: 'sqlite',
+			filename: normaliseSQLiteUrl(credentials.url, 'bun'),
+		});
+
+		const db = drizzle({ client });
+		const migrateFn = async (config: MigrationConfig) => {
+			return migrate(db, config);
+		};
+
+		const query = async (sql: string, params?: any[]) => {
+			const result = await client.unsafe(sql, params ?? []);
+			return result;
+		};
+		const run = async (sql: string) => {
+			await client.unsafe(sql);
+		};
+
+		const proxy: Proxy = async (params) => {
+			const query = client.unsafe(params.sql, params.params);
+			if (params.mode === 'array') {
+				return await query.values();
+			}
+			return await query;
+		};
+
+		const transactionProxy: TransactionProxy = async (queries) => {
+			const results: any[] = [];
+			try {
+				await client.transaction(async (tx) => {
+					for (const query of queries) {
+						const result = await tx.unsafe(query.sql);
+						results.push(result);
+					}
+				});
+			} catch (error) {
+				results.push(error as Error);
+			}
+			return results;
+		};
+
+		return {
+			packageName: 'bun',
+			query,
+			run,
+			proxy,
+			transactionProxy,
+			migrate: migrateFn,
+		};
+	}
+
 	console.log(
-		"Please install either 'better-sqlite3' or '@libsql/client' for Drizzle Kit to connect to SQLite databases",
+		"Please install either 'better-sqlite3', '@libsql/client' or 'bun' for Drizzle Kit to connect to SQLite databases",
 	);
 	process.exit(1);
 };

drizzle-kit/src/utils/utils-node.ts

@@ -300,7 +300,7 @@ export const validateWithReport = (snapshots: string[], dialect: Dialect) => {
 
 export const normaliseSQLiteUrl = (
 	it: string,
-	type: 'libsql' | 'better-sqlite' | '@tursodatabase/database',
+	type: 'libsql' | 'better-sqlite' | '@tursodatabase/database' | 'bun',
 ) => {
 	if (type === 'libsql') {
 		if (it.startsWith('file:')) {
@@ -317,7 +317,7 @@ export const normaliseSQLiteUrl = (
 		}
 	}
 
-	if (type === 'better-sqlite' || type === '@tursodatabase/database') {
+	if (type === 'better-sqlite' || type === '@tursodatabase/database' || type === 'bun') {
 		if (it.startsWith('file:')) {
 			return it.substring(5);
 		}

drizzle-orm/src/cockroach-core/dialect.ts

@@ -252,6 +252,23 @@ export class CockroachDialect {
 					} else {
 						chunk.push(field.isAlias ? sql`${getOriginalColumnFromAlias(field)} as ${field}` : field);
 					}
+				} else if (is(field, Subquery)) {
+					const entries = Object.entries(field._.selectedFields) as [string, SQL.Aliased | Column | SQL][];
+
+					if (entries.length === 1) {
+						const entry = entries[0]![1];
+
+						const fieldDecoder = is(entry, SQL)
+							? entry.decoder
+							: is(entry, Column)
+							? { mapFromDriverValue: (v: any) => entry.mapFromDriverValue(v) }
+							: entry.sql.decoder;
+
+						if (fieldDecoder) {
+							field._.sql.decoder = fieldDecoder;
+						}
+					}
+					chunk.push(field);
 				}
 
 				if (i < columnsLen - 1) {