feat(db0): add db0 driver for drizzle-orm

Author: onmax

drizzle-orm/package.json

@@ -72,7 +72,8 @@
 		"sql.js": ">=1",
 		"sqlite3": ">=5",
 		"gel": ">=2",
-		"@upstash/redis": ">=1.34.7"
+		"@upstash/redis": ">=1.34.7",
+		"db0": ">=0.2.0"
 	},
 	"peerDependenciesMeta": {
 		"mysql2": {
@@ -161,6 +162,9 @@
 		},
 		"@upstash/redis": {
 			"optional": true
+		},
+		"db0": {
+			"optional": true
 		}
 	},
 	"devDependencies": {
@@ -188,6 +192,7 @@
 		"better-sqlite3": "^11.9.1",
 		"bun-types": "^1.2.0",
 		"cpy": "^10.1.0",
+		"db0": "^0.2.3",
 		"expo-sqlite": "^14.0.0",
 		"gel": "^2.0.0",
 		"glob": "^11.0.1",

drizzle-orm/src/db0/driver.ts

@@ -0,0 +1,48 @@
+import type { Database } from 'db0';
+import type { DrizzleConfig } from '~/utils.ts';
+import { isConfig } from '~/utils.ts';
+import { constructPg, type Db0PgDatabase } from './pg/index.ts';
+import { constructSqlite, type Db0SQLiteDatabase } from './sqlite/index.ts';
+
+export type Db0Database = Db0SQLiteDatabase<any> | Db0PgDatabase<any>;
+
+export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
+	...params:
+		| [Database]
+		| [Database, DrizzleConfig<TSchema>]
+		| [DrizzleConfig<TSchema> & { client: Database }]
+): Db0Database & { $client: Database } {
+	if (isConfig(params[0])) {
+		const { client, ...drizzleConfig } = params[0] as DrizzleConfig<TSchema> & { client: Database };
+		return drizzleInternal(client, drizzleConfig);
+	}
+
+	return drizzleInternal(params[0] as Database, params[1] as DrizzleConfig<TSchema> | undefined);
+}
+
+function drizzleInternal<TSchema extends Record<string, unknown> = Record<string, never>>(
+	client: Database,
+	config: DrizzleConfig<TSchema> = {},
+): Db0Database & { $client: Database } {
+	const dialect = client.dialect;
+
+	switch (dialect) {
+		case 'sqlite':
+		case 'libsql':
+			return constructSqlite(client, config) as any;
+		case 'postgresql':
+			return constructPg(client, config) as any;
+		case 'mysql':
+			throw new Error('drizzle-orm/db0: MySQL support is not yet implemented');
+		default:
+			throw new Error(`drizzle-orm/db0: Unsupported db0 dialect: ${dialect}`);
+	}
+}
+
+export namespace drizzle {
+	export function mock<TSchema extends Record<string, unknown> = Record<string, never>>(
+		config?: DrizzleConfig<TSchema>,
+	): Db0Database & { $client: '$client is not available on drizzle.mock()' } {
+		return constructSqlite({} as any, config) as any;
+	}
+}

drizzle-orm/src/db0/index.ts

@@ -0,0 +1,20 @@
+export { drizzle, type Db0Database } from './driver.ts';
+export {
+	constructPg,
+	Db0PgDatabase,
+	Db0PgPreparedQuery,
+	type Db0PgQueryResult,
+	type Db0PgQueryResultHKT,
+	Db0PgSession,
+	type Db0PgSessionOptions,
+	Db0PgTransaction,
+} from './pg/index.ts';
+export {
+	constructSqlite,
+	Db0SQLiteDatabase,
+	Db0SQLitePreparedQuery,
+	Db0SQLiteSession,
+	Db0SQLiteTransaction,
+	type Db0RunResult,
+	type Db0SQLiteSessionOptions,
+} from './sqlite/index.ts';

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

@@ -0,0 +1,49 @@
+import type { Database } from 'db0';
+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 } from '~/utils.ts';
+import { type Db0PgQueryResultHKT, Db0PgSession, type Db0PgSessionOptions } from './session.ts';
+
+export class Db0PgDatabase<
+	TSchema extends Record<string, unknown> = Record<string, never>,
+> extends PgDatabase<Db0PgQueryResultHKT, TSchema> {
+	static override readonly [entityKind]: string = 'Db0PgDatabase';
+}
+
+export function constructPg<TSchema extends Record<string, unknown> = Record<string, never>>(
+	client: Database,
+	config: DrizzleConfig<TSchema> = {},
+): Db0PgDatabase<TSchema> & { $client: Database } {
+	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 sessionOptions: Db0PgSessionOptions = { logger, cache: config.cache };
+	const session = new Db0PgSession(client, dialect, schema, sessionOptions);
+	const db = new Db0PgDatabase(dialect, session, schema as any) as Db0PgDatabase<TSchema>;
+	(<any>db).$client = client;
+
+	return db as any;
+}

drizzle-orm/src/db0/pg/index.ts

@@ -0,0 +1,10 @@
+export { constructPg, Db0PgDatabase } from './driver.ts';
+export * from './migrator.ts';
+export {
+	Db0PgPreparedQuery,
+	type Db0PgQueryResult,
+	type Db0PgQueryResultHKT,
+	Db0PgSession,
+	type Db0PgSessionOptions,
+	Db0PgTransaction,
+} from './session.ts';

drizzle-orm/src/db0/pg/migrator.ts

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

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

@@ -0,0 +1,197 @@
+import type { Database, Primitive } from 'db0';
+import { type Cache, NoopCache } from '~/cache/core/cache.ts';
+import type { WithCacheConfig } from '~/cache/core/types.ts';
+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, type SQL, sql } from '~/sql/sql.ts';
+import { mapResultRow } from '~/utils.ts';
+
+export interface Db0PgSessionOptions {
+	logger?: Logger;
+	cache?: Cache;
+}
+
+export interface Db0PgQueryResult {
+	rows: unknown[];
+	rowCount: number;
+}
+
+export class Db0PgPreparedQuery<T extends PreparedQueryConfig = PreparedQueryConfig> extends PgPreparedQuery<T> {
+	static override readonly [entityKind]: string = 'Db0PgPreparedQuery';
+
+	constructor(
+		private client: Database,
+		private queryString: string,
+		private params: unknown[],
+		private logger: Logger,
+		cache: Cache,
+		queryMetadata: { type: 'select' | 'update' | 'delete' | 'insert'; tables: string[] } | undefined,
+		cacheConfig: WithCacheConfig | undefined,
+		private fields: SelectedFieldsOrdered | undefined,
+		name: string | undefined,
+		private _isResponseInArrayMode: boolean,
+		private customResultMapper?: (rows: unknown[][], mapColumnValue?: (value: unknown) => unknown) => T['execute'],
+	) {
+		super({ sql: queryString, params }, cache, queryMetadata, cacheConfig);
+	}
+
+	async execute(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['execute']> {
+		const params = fillPlaceholders(this.params, placeholderValues) as Primitive[];
+		this.logger.logQuery(this.queryString, params);
+
+		const { fields, client, customResultMapper, queryString, joinsNotNullableMap } = this;
+
+		if (!fields && !customResultMapper) {
+			return await this.queryWithCache(queryString, params, async () => {
+				const stmt = client.prepare(queryString);
+				const rows = await stmt.all(...params) as unknown[];
+				return { rows, rowCount: rows.length } as T['execute'];
+			});
+		}
+
+		// db0 doesn't have array mode, so we get objects and convert to arrays
+		return await this.queryWithCache(queryString, params, async () => {
+			const stmt = client.prepare(queryString);
+			const rows = await stmt.all(...params) as Record<string, unknown>[];
+			const arrayRows = rows.map((row) => Object.values(row));
+
+			if (customResultMapper) {
+				return customResultMapper(arrayRows);
+			}
+
+			return arrayRows.map((row) => mapResultRow<T['execute']>(fields!, row, joinsNotNullableMap));
+		});
+	}
+
+	async all(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['all']> {
+		const params = fillPlaceholders(this.params, placeholderValues) as Primitive[];
+		this.logger.logQuery(this.queryString, params);
+		return await this.queryWithCache(this.queryString, params, async () => {
+			const stmt = this.client.prepare(this.queryString);
+			return stmt.all(...params) as Promise<T['all']>;
+		});
+	}
+
+	/** @internal */
+	isResponseInArrayMode(): boolean {
+		return this._isResponseInArrayMode;
+	}
+}
+
+export class Db0PgSession<
+	TFullSchema extends Record<string, unknown>,
+	TSchema extends TablesRelationalConfig,
+> extends PgSession<Db0PgQueryResultHKT, TFullSchema, TSchema> {
+	static override readonly [entityKind]: string = 'Db0PgSession';
+
+	private logger: Logger;
+	private cache: Cache;
+
+	constructor(
+		private client: Database,
+		dialect: PgDialect,
+		private schema: RelationalSchemaConfig<TSchema> | undefined,
+		private options: Db0PgSessionOptions = {},
+	) {
+		super(dialect);
+		this.logger = options.logger ?? new NoopLogger();
+		this.cache = options.cache ?? new NoopCache();
+	}
+
+	prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
+		query: Query,
+		fields: SelectedFieldsOrdered | undefined,
+		name: string | undefined,
+		isResponseInArrayMode: boolean,
+		customResultMapper?: (rows: unknown[][], mapColumnValue?: (value: unknown) => unknown) => T['execute'],
+		queryMetadata?: { type: 'select' | 'update' | 'delete' | 'insert'; tables: string[] },
+		cacheConfig?: WithCacheConfig,
+	): PgPreparedQuery<T> {
+		return new Db0PgPreparedQuery(
+			this.client,
+			query.sql,
+			query.params,
+			this.logger,
+			this.cache,
+			queryMetadata,
+			cacheConfig,
+			fields,
+			name,
+			isResponseInArrayMode,
+			customResultMapper,
+		);
+	}
+
+	override async transaction<T>(
+		transaction: (tx: Db0PgTransaction<TFullSchema, TSchema>) => Promise<T>,
+		config?: PgTransactionConfig,
+	): Promise<T> {
+		const tx = new Db0PgTransaction<TFullSchema, TSchema>(this.dialect, this, this.schema);
+
+		let beginSql = 'begin';
+		if (config) {
+			const chunks: string[] = [];
+			if (config.isolationLevel) {
+				chunks.push(`isolation level ${config.isolationLevel}`);
+			}
+			if (config.accessMode) {
+				chunks.push(config.accessMode);
+			}
+			if (typeof config.deferrable === 'boolean') {
+				chunks.push(config.deferrable ? 'deferrable' : 'not deferrable');
+			}
+			if (chunks.length > 0) {
+				beginSql = `begin ${chunks.join(' ')}`;
+			}
+		}
+
+		await this.execute(sql.raw(beginSql));
+		try {
+			const result = await transaction(tx);
+			await this.execute(sql`commit`);
+			return result;
+		} catch (err) {
+			await this.execute(sql`rollback`);
+			throw err;
+		}
+	}
+
+	override async count(countSql: SQL): Promise<number> {
+		const res = await this.execute<Db0PgQueryResult>(countSql);
+		return Number((res.rows[0] as Record<string, unknown>)['count']);
+	}
+}
+
+export class Db0PgTransaction<
+	TFullSchema extends Record<string, unknown>,
+	TSchema extends TablesRelationalConfig,
+> extends PgTransaction<Db0PgQueryResultHKT, TFullSchema, TSchema> {
+	static override readonly [entityKind]: string = 'Db0PgTransaction';
+
+	override async transaction<T>(
+		transaction: (tx: Db0PgTransaction<TFullSchema, TSchema>) => Promise<T>,
+	): Promise<T> {
+		const savepointName = `sp${this.nestedIndex + 1}`;
+		const tx = new Db0PgTransaction<TFullSchema, TSchema>(this.dialect, this.session, this.schema, this.nestedIndex + 1);
+		await tx.execute(sql.raw(`savepoint ${savepointName}`));
+		try {
+			const result = await transaction(tx);
+			await tx.execute(sql.raw(`release savepoint ${savepointName}`));
+			return result;
+		} catch (err) {
+			await tx.execute(sql.raw(`rollback to savepoint ${savepointName}`));
+			throw err;
+		}
+	}
+}
+
+export interface Db0PgQueryResultHKT extends PgQueryResultHKT {
+	type: Db0PgQueryResult;
+}