Add bun sql driver

drizzle-orm/package.json

@@ -185,7 +185,7 @@
 		"@vercel/postgres": "^0.8.0",
 		"@xata.io/client": "^0.29.3",
 		"better-sqlite3": "^8.4.0",
-		"bun-types": "^0.6.6",
+		"bun-types": "^1.2.0",
 		"cpy": "^10.1.0",
 		"expo-sqlite": "^14.0.0",
 		"knex": "^2.4.2",

drizzle-orm/src/bun-sql/driver.ts

@@ -0,0 +1,122 @@
+/// <reference types="bun-types" />
+
+import type { SQLOptions } from 'bun';
+import { SQL } from 'bun';
+import { entityKind } from '~/entity.ts';
+import { DefaultLogger } from '~/logger.ts';
+import { PgDatabase } from '~/pg-core/db.ts';
+import { PgDialect } from '~/pg-core/dialect.ts';
+import {
+	createTableRelationsHelpers,
+	extractTablesRelationalConfig,
+	type RelationalSchemaConfig,
+	type TablesRelationalConfig,
+} from '~/relations.ts';
+import { type DrizzleConfig, isConfig } from '~/utils.ts';
+import type { BunSQLQueryResultHKT } from './session.ts';
+import { BunSQLSession } from './session.ts';
+
+export class BunSQLDatabase<
+	TSchema extends Record<string, unknown> = Record<string, never>,
+> extends PgDatabase<BunSQLQueryResultHKT, TSchema> {
+	static override readonly [entityKind]: string = 'BunSQLDatabase';
+}
+
+function construct<TSchema extends Record<string, unknown> = Record<string, never>>(
+	client: SQL,
+	config: DrizzleConfig<TSchema> = {},
+): BunSQLDatabase<TSchema> & {
+	$client: SQL;
+} {
+	const dialect = new PgDialect({ casing: config.casing });
+	let logger;
+	if (config.logger === true) {
+		logger = new DefaultLogger();
+	} else if (config.logger !== false) {
+		logger = config.logger;
+	}
+
+	let schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined;
+	if (config.schema) {
+		const tablesConfig = extractTablesRelationalConfig(
+			config.schema,
+			createTableRelationsHelpers,
+		);
+		schema = {
+			fullSchema: config.schema,
+			schema: tablesConfig.tables,
+			tableNamesMap: tablesConfig.tableNamesMap,
+		};
+	}
+
+	const session = new BunSQLSession(client, dialect, schema, { logger });
+	const db = new BunSQLDatabase(dialect, session, schema as any) as BunSQLDatabase<TSchema>;
+	(<any> db).$client = client;
+
+	return db as any;
+}
+
+export function drizzle<
+	TSchema extends Record<string, unknown> = Record<string, never>,
+	TClient extends SQL = SQL,
+>(
+	...params: [
+		TClient | string,
+	] | [
+		TClient | string,
+		DrizzleConfig<TSchema>,
+	] | [
+		(
+			& DrizzleConfig<TSchema>
+			& ({
+				connection: string | ({ url?: string } & SQLOptions);
+			} | {
+				client: TClient;
+			})
+		),
+	]
+): BunSQLDatabase<TSchema> & {
+	$client: TClient;
+} {
+	if (typeof params[0] === 'string') {
+		const instance = new SQL(params[0]);
+
+		return construct(instance, params[1]) as any;
+	}
+
+	if (isConfig(params[0])) {
+		const { connection, client, ...drizzleConfig } = params[0] as {
+			connection?: { url?: string } & SQLOptions;
+			client?: TClient;
+		} & DrizzleConfig<TSchema>;
+
+		if (client) return construct(client, drizzleConfig) as any;
+
+		if (typeof connection === 'object' && connection.url !== undefined) {
+			const { url, ...config } = connection;
+
+			const instance = new SQL({ url, ...config });
+			return construct(instance, drizzleConfig) as any;
+		}
+
+		const instance = new SQL(connection);
+		return construct(instance, drizzleConfig) as any;
+	}
+
+	return construct(params[0] as TClient, params[1] as DrizzleConfig<TSchema> | undefined) as any;
+}
+
+export namespace drizzle {
+	export function mock<TSchema extends Record<string, unknown> = Record<string, never>>(
+		config?: DrizzleConfig<TSchema>,
+	): BunSQLDatabase<TSchema> & {
+		$client: '$client is not available on drizzle.mock()';
+	} {
+		return construct({
+			options: {
+				parsers: {},
+				serializers: {},
+			},
+		} as any, config) as any;
+	}
+}

drizzle-orm/src/bun-sql/index.ts

@@ -0,0 +1,2 @@
+export * from './driver.ts';
+export * from './session.ts';

drizzle-orm/src/bun-sql/migrator.ts

@@ -0,0 +1,11 @@
+import type { MigrationConfig } from '~/migrator.ts';
+import { readMigrationFiles } from '~/migrator.ts';
+import type { BunSQLDatabase } from './driver.ts';
+
+export async function migrate<TSchema extends Record<string, unknown>>(
+	db: BunSQLDatabase<TSchema>,
+	config: MigrationConfig,
+) {
+	const migrations = readMigrationFiles(config);
+	await db.dialect.migrate(migrations, db.session, config);
+}

drizzle-orm/src/bun-sql/session.ts

@@ -0,0 +1,199 @@
+/// <reference types="bun-types" />
+
+import type { SavepointSQL, SQL, TransactionSQL } from 'bun';
+import { entityKind } from '~/entity.ts';
+import type { Logger } from '~/logger.ts';
+import { NoopLogger } from '~/logger.ts';
+import type { PgDialect } from '~/pg-core/dialect.ts';
+import { PgTransaction } from '~/pg-core/index.ts';
+import type { SelectedFieldsOrdered } from '~/pg-core/query-builders/select.types.ts';
+import type { PgQueryResultHKT, PgTransactionConfig, PreparedQueryConfig } from '~/pg-core/session.ts';
+import { PgPreparedQuery, PgSession } from '~/pg-core/session.ts';
+import type { RelationalSchemaConfig, TablesRelationalConfig } from '~/relations.ts';
+import { fillPlaceholders, type Query } from '~/sql/sql.ts';
+import { tracer } from '~/tracing.ts';
+import { type Assume, mapResultRow } from '~/utils.ts';
+
+export class BunSQLPreparedQuery<T extends PreparedQueryConfig> extends PgPreparedQuery<T> {
+	static override readonly [entityKind]: string = 'BunSQLPreparedQuery';
+
+	constructor(
+		private client: SQL,
+		private queryString: string,
+		private params: unknown[],
+		private logger: Logger,
+		private fields: SelectedFieldsOrdered | undefined,
+		private _isResponseInArrayMode: boolean,
+		private customResultMapper?: (rows: unknown[][]) => T['execute'],
+	) {
+		super({ sql: queryString, params });
+	}
+
+	async execute(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['execute']> {
+		return tracer.startActiveSpan('drizzle.execute', async (span) => {
+			const params = fillPlaceholders(this.params, placeholderValues);
+
+			span?.setAttributes({
+				'drizzle.query.text': this.queryString,
+				'drizzle.query.params': JSON.stringify(params),
+			});
+
+			this.logger.logQuery(this.queryString, params);
+
+			const { fields, queryString: query, client, joinsNotNullableMap, customResultMapper } = this;
+			if (!fields && !customResultMapper) {
+				return tracer.startActiveSpan('drizzle.driver.execute', () => {
+					return client.unsafe(query, params as any[]);
+				});
+			}
+
+			const rows: any[] = await tracer.startActiveSpan('drizzle.driver.execute', () => {
+				span?.setAttributes({
+					'drizzle.query.text': query,
+					'drizzle.query.params': JSON.stringify(params),
+				});
+
+				return client.unsafe(query, params as any[]).values();
+			});
+
+			return tracer.startActiveSpan('drizzle.mapResponse', () => {
+				return customResultMapper
+					? customResultMapper(rows)
+					: rows.map((row) => mapResultRow<T['execute']>(fields!, row, joinsNotNullableMap));
+			});
+		});
+	}
+
+	all(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['all']> {
+		return tracer.startActiveSpan('drizzle.execute', async (span) => {
+			const params = fillPlaceholders(this.params, placeholderValues);
+			span?.setAttributes({
+				'drizzle.query.text': this.queryString,
+				'drizzle.query.params': JSON.stringify(params),
+			});
+			this.logger.logQuery(this.queryString, params);
+			return tracer.startActiveSpan('drizzle.driver.execute', () => {
+				span?.setAttributes({
+					'drizzle.query.text': this.queryString,
+					'drizzle.query.params': JSON.stringify(params),
+				});
+				return this.client.unsafe(this.queryString, params as any[]);
+			});
+		});
+	}
+
+	/** @internal */
+	isResponseInArrayMode(): boolean {
+		return this._isResponseInArrayMode;
+	}
+}
+
+export interface BunSQLSessionOptions {
+	logger?: Logger;
+}
+
+export class BunSQLSession<
+	TSQL extends SQL,
+	TFullSchema extends Record<string, unknown>,
+	TSchema extends TablesRelationalConfig,
+> extends PgSession<BunSQLQueryResultHKT, TFullSchema, TSchema> {
+	static override readonly [entityKind]: string = 'BunSQLSession';
+
+	logger: Logger;
+
+	constructor(
+		public client: TSQL,
+		dialect: PgDialect,
+		private schema: RelationalSchemaConfig<TSchema> | undefined,
+		/** @internal */
+		readonly options: BunSQLSessionOptions = {},
+	) {
+		super(dialect);
+		this.logger = options.logger ?? new NoopLogger();
+	}
+
+	prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
+		query: Query,
+		fields: SelectedFieldsOrdered | undefined,
+		name: string | undefined,
+		isResponseInArrayMode: boolean,
+		customResultMapper?: (rows: unknown[][]) => T['execute'],
+	): PgPreparedQuery<T> {
+		return new BunSQLPreparedQuery(
+			this.client,
+			query.sql,
+			query.params,
+			this.logger,
+			fields,
+			isResponseInArrayMode,
+			customResultMapper,
+		);
+	}
+
+	query(query: string, params: unknown[]): Promise<any> {
+		this.logger.logQuery(query, params);
+		return this.client.unsafe(query, params as any[]).values();
+	}
+
+	queryObjects(
+		query: string,
+		params: unknown[],
+	): Promise<any> {
+		return this.client.unsafe(query, params as any[]);
+	}
+
+	override transaction<T>(
+		transaction: (tx: BunSQLTransaction<TFullSchema, TSchema>) => Promise<T>,
+		config?: PgTransactionConfig,
+	): Promise<T> {
+		return this.client.begin(async (client) => {
+			const session = new BunSQLSession<TransactionSQL, TFullSchema, TSchema>(
+				client,
+				this.dialect,
+				this.schema,
+				this.options,
+			);
+			const tx = new BunSQLTransaction(this.dialect, session, this.schema);
+			if (config) {
+				await tx.setTransaction(config);
+			}
+			return transaction(tx);
+		}) as Promise<T>;
+	}
+}
+
+export class BunSQLTransaction<
+	TFullSchema extends Record<string, unknown>,
+	TSchema extends TablesRelationalConfig,
+> extends PgTransaction<BunSQLQueryResultHKT, TFullSchema, TSchema> {
+	static override readonly [entityKind]: string = 'BunSQLTransaction';
+
+	constructor(
+		dialect: PgDialect,
+		/** @internal */
+		override readonly session: BunSQLSession<TransactionSQL | SavepointSQL, TFullSchema, TSchema>,
+		schema: RelationalSchemaConfig<TSchema> | undefined,
+		nestedIndex = 0,
+	) {
+		super(dialect, session, schema, nestedIndex);
+	}
+
+	override transaction<T>(
+		transaction: (tx: BunSQLTransaction<TFullSchema, TSchema>) => Promise<T>,
+	): Promise<T> {
+		return (this.session.client as TransactionSQL).savepoint((client) => {
+			const session = new BunSQLSession<SavepointSQL, TFullSchema, TSchema>(
+				client,
+				this.dialect,
+				this.schema,
+				this.session.options,
+			);
+			const tx = new BunSQLTransaction<TFullSchema, TSchema>(this.dialect, session, this.schema);
+			return transaction(tx);
+		}) as Promise<T>;
+	}
+}
+
+export interface BunSQLQueryResultHKT extends PgQueryResultHKT {
+	type: Assume<this['row'], Record<string, any>[]>;
+}

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

@@ -117,7 +117,7 @@ export class PreparedQuery<T extends PreparedQueryConfig = PreparedQueryConfig>
 		super('sync', executeMethod, query);
 	}
 
-	run(placeholderValues?: Record<string, unknown>): void {
+	run(placeholderValues?: Record<string, unknown>) {
 		const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
 		this.logger.logQuery(this.query.sql, params);
 		return this.stmt.run(...params);