Edge and Mobile SQLite

Copyexport function drizzle<
	TSchema extends Record<string, unknown> = Record<string, never>,
	TClient extends AnyD1Database = AnyD1Database,
>(
	client: TClient,
	config: DrizzleConfig<TSchema> = {},
): DrizzleD1Database<TSchema> & {
	$client: TClient;
} {
	const dialect = new SQLiteAsyncDialect({ 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 SQLiteD1Session(client as D1Database, dialect, schema, { logger, cache: config.cache });
	const db = new DrizzleD1Database('async', dialect, session, schema) as DrizzleD1Database<TSchema>;
	(<any> db).$client = client;
	(<any> db).$cache = config.cache;
	if ((<any> db).$cache) {
		(<any> db).$cache['invalidate'] = config.cache?.onMutate;
	}
 
	return db as any;
}

Copyconst session = new SQLiteD1Session(client as D1Database, dialect, schema, { logger, cache: config.cache });
const db = new DrizzleD1Database('async', dialect, session, schema) as DrizzleD1Database<TSchema>;
(<any> db).$client = client;
(<any> db).$cache = config.cache;
if ((<any> db).$cache) {
	(<any> db).$cache['invalidate'] = config.cache?.onMutate;
}

CopyprepareQuery(
	query: Query,
	fields: SelectedFieldsOrdered | undefined,
	executeMethod: SQLiteExecuteMethod,
	isResponseInArrayMode: boolean,
	customResultMapper?: (rows: unknown[][]) => unknown,
	queryMetadata?: {
		type: 'select' | 'update' | 'delete' | 'insert';
		tables: string[];
	},
	cacheConfig?: WithCacheConfig,
): D1PreparedQuery {
	const stmt = this.client.prepare(query.sql);
	return new D1PreparedQuery(
		stmt,
		query,
		this.logger,
		this.cache,
		queryMetadata,
		cacheConfig,
		fields,
		executeMethod,
		isResponseInArrayMode,
		customResultMapper,
	);
}

Copyasync batch<T extends BatchItem<'sqlite'>[] | readonly BatchItem<'sqlite'>[]>(queries: T) {
	const preparedQueries: PreparedQuery[] = [];
	const builtQueries: D1PreparedStatement[] = [];
 
	for (const query of queries) {
		const preparedQuery = query._prepare();
		const builtQuery = preparedQuery.getQuery();
		preparedQueries.push(preparedQuery);
		if (builtQuery.params.length > 0) {
			builtQueries.push((preparedQuery as D1PreparedQuery).stmt.bind(...builtQuery.params));
		} else {
			const builtQuery = preparedQuery.getQuery();
			builtQueries.push(
				this.client.prepare(builtQuery.sql).bind(...builtQuery.params),
			);
		}
	}
 
	const batchResults = await this.client.batch<any>(builtQueries);
	return batchResults.map((result, i) => preparedQueries[i]!.mapResult(result, true));
}

Copyoverride async transaction<T>(
	transaction: (tx: D1Transaction<TFullSchema, TSchema>) => T | Promise<T>,
	config?: SQLiteTransactionConfig,
): Promise<T> {
	const tx = new D1Transaction('async', this.dialect, this, this.schema);
	await this.run(sql.raw(`begin${config?.behavior ? ' ' + config.behavior : ''}`));
	try {
		const result = await transaction(tx);
		await this.run(sql`commit`);
		return result;
	} catch (err) {
		await this.run(sql`rollback`);
		throw err;
	}
}

Copyoverride async transaction<T>(transaction: (tx: D1Transaction<TFullSchema, TSchema>) => Promise<T>): Promise<T> {
	const savepointName = `sp${this.nestedIndex}`;
	const tx = new D1Transaction('async', this.dialect, this.session, this.schema, this.nestedIndex + 1);
	await this.session.run(sql.raw(`savepoint ${savepointName}`));
	try {
		const result = await transaction(tx);
		await this.session.run(sql.raw(`release savepoint ${savepointName}`));
		return result;
	} catch (err) {
		await this.session.run(sql.raw(`rollback to savepoint ${savepointName}`));
		throw err;
	}
}

Copyexport function drizzle<
	TSchema extends Record<string, unknown> = Record<string, never>,
	TClient extends DurableObjectStorage = DurableObjectStorage,
>(
	client: TClient,
	config: DrizzleConfig<TSchema> = {},
): DrizzleSqliteDODatabase<TSchema> & {
	$client: TClient;
} {
	const dialect = new SQLiteSyncDialect({ 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,
		};
	}

Copy	const session = new SQLiteDOSession(client as DurableObjectStorage, dialect, schema, { logger });
	const db = new DrizzleSqliteDODatabase('sync', dialect, session, schema) as DrizzleSqliteDODatabase<TSchema>;
	(<any> db).$client = client;
 
	return db as any;
}

Copyexport class SQLiteDOSession<TFullSchema extends Record<string, unknown>, TSchema extends TablesRelationalConfig>
	extends SQLiteSession<
		'sync',
		SqlStorageCursor<Record<string, SqlStorageValue>>,
		TFullSchema,
		TSchema
	>
{
	static override readonly [entityKind]: string = 'SQLiteDOSession';
 
	private logger: Logger;
 
	constructor(
		private client: DurableObjectStorage,
		dialect: SQLiteSyncDialect,
		private schema: RelationalSchemaConfig<TSchema> | undefined,
		options: SQLiteDOSessionOptions = {},
	) {
		super(dialect);
		this.logger = options.logger ?? new NoopLogger();
	}
 
	prepareQuery<T extends Omit<PreparedQueryConfig, 'run'>>(
		query: Query,
		fields: SelectedFieldsOrdered | undefined,
		executeMethod: SQLiteExecuteMethod,
		isResponseInArrayMode: boolean,
		customResultMapper?: (rows: unknown[][]) => unknown,
	): SQLiteDOPreparedQuery<T> {
		return new SQLiteDOPreparedQuery(
			this.client,
			query,
			this.logger,
			fields,
			executeMethod,
			isResponseInArrayMode,
			customResultMapper,
		);
	}
 
	override transaction<T>(
		transaction: (
			tx: SQLiteTransaction<'sync', SqlStorageCursor<Record<string, SqlStorageValue>>, TFullSchema, TSchema>,
		) => T,
		_config?: SQLiteTransactionConfig,
	): T {
		const tx = new SQLiteDOTransaction('sync', this.dialect, this, this.schema);
		return this.client.transactionSync(() => transaction(tx));
	}
}

Copyexport class SQLiteDOPreparedQuery<T extends PreparedQueryConfig = PreparedQueryConfig> extends PreparedQueryBase<{
	type: 'sync';
	run: void;
	all: T['all'];
	get: T['get'];
	values: T['values'];
	execute: T['execute'];
}> {
	static override readonly [entityKind]: string = 'SQLiteDOPreparedQuery';
 
	constructor(
		private client: DurableObjectStorage,
		query: Query,
		private logger: Logger,
		private fields: SelectedFieldsOrdered | undefined,
		executeMethod: SQLiteExecuteMethod,
		private _isResponseInArrayMode: boolean,
		private customResultMapper?: (rows: unknown[][]) => unknown,
	) {
		// 3-6 params are for cache. As long as we don't support sync cache - it will be skipped here
		super('sync', executeMethod, query, undefined, undefined, undefined);
	}
 
	run(placeholderValues?: Record<string, unknown>): void {
		const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
		this.logger.logQuery(this.query.sql, params);
 
		params.length > 0 ? this.client.sql.exec(this.query.sql, ...params) : this.client.sql.exec(this.query.sql);
	}
 
	all(placeholderValues?: Record<string, unknown>): T['all'] {
		const { fields, joinsNotNullableMap, query, logger, client, customResultMapper } = this;
		if (!fields && !customResultMapper) {
			const params = fillPlaceholders(query.params, placeholderValues ?? {});
			logger.logQuery(query.sql, params);
 
			return params.length > 0 ? client.sql.exec(query.sql, ...params).toArray() : client.sql.exec(query.sql).toArray();
		}
 
		const rows = this.values(placeholderValues) as unknown[][];
 
		if (customResultMapper) {
			return customResultMapper(rows) as T['all'];
		}
 
		return rows.map((row) => mapResultRow(fields!, row, joinsNotNullableMap));
	}
 
	get(placeholderValues?: Record<string, unknown>): T['get'] {
		const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
		this.logger.logQuery(this.query.sql, params);
 
		const { fields, client, joinsNotNullableMap, customResultMapper, query } = this;
		if (!fields && !customResultMapper) {
			return (params.length > 0 ? client.sql.exec(query.sql, ...params) : client.sql.exec(query.sql)).next().value;
		}
 
		const rows = this.values(placeholderValues) as unknown[][];
		const row = rows[0];
 
		if (!row) {
			return undefined;
		}
 
		if (customResultMapper) {
			return customResultMapper(rows) as T['get'];
		}
 
		return mapResultRow(fields!, row, joinsNotNullableMap);
	}
 
	values(placeholderValues?: Record<string, unknown>): T['values'] {
		const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
		this.logger.logQuery(this.query.sql, params);
 
		const res = params.length > 0
			? this.client.sql.exec(this.query.sql, ...params)
			: this.client.sql.exec(this.query.sql);
 
		// @ts-ignore .raw().toArray() exists
		return res.raw().toArray();
	}
}

Copyexport function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
	client: SQLiteDatabase,
	config: DrizzleConfig<TSchema> = {},
): ExpoSQLiteDatabase<TSchema> & {
	$client: SQLiteDatabase;
} {
	const dialect = new SQLiteSyncDialect({ 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,
		};
	}

Copy	const session = new ExpoSQLiteSession(client, dialect, schema, { logger });
	const db = new ExpoSQLiteDatabase('sync', dialect, session, schema) as ExpoSQLiteDatabase<TSchema>;
	(<any> db).$client = client;
 
	return db as any;
}

Copyexport function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
	client: OPSQLiteConnection,
	config: DrizzleConfig<TSchema> = {},
): OPSQLiteDatabase<TSchema> & {
	$client: OPSQLiteConnection;
} {
	const dialect = new SQLiteAsyncDialect({ casing: config.casing });

Copylet 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,
	};
}

Copyconst session = new OPSQLiteSession(client, dialect, schema, { logger, cache: config.cache });
const db = new OPSQLiteDatabase('async', dialect, session, schema) as OPSQLiteDatabase<TSchema>;
(<any> db).$client = client;
(<any> db).$cache = config.cache;
if ((<any> db).$cache) {
	(<any> db).$cache['invalidate'] = config.cache?.onMutate;
}
 
return db as any;

Copyexport function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
	client: Database,
	config: DrizzleConfig<TSchema> = {},
): SQLJsDatabase<TSchema> {

Copyconst dialect = new SQLiteSyncDialect({ 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,
	};
}

Copyconst session = new SQLJsSession(client, dialect, schema, { logger });
return new BaseSQLiteDatabase('sync', dialect, session, schema) as SQLJsDatabase<TSchema>;

Copy	override transaction<T>(
		transaction: (tx: SQLJsTransaction<TFullSchema, TSchema>) => T,
		config: SQLiteTransactionConfig = {},
	): T {
		const tx = new SQLJsTransaction('sync', this.dialect, this, this.schema);
		this.run(sql.raw(`begin${config.behavior ? ` ${config.behavior}` : ''}`));
		try {
			const result = transaction(tx);
			this.run(sql`commit`);
			return result;
		} catch (err) {
			this.run(sql`rollback`);
			throw err;
		}
	}
}
 
export class SQLJsTransaction<
	TFullSchema extends Record<string, unknown>,
	TSchema extends TablesRelationalConfig,
> extends SQLiteTransaction<'sync', void, TFullSchema, TSchema> {
	static override readonly [entityKind]: string = 'SQLJsTransaction';
 
	override transaction<T>(transaction: (tx: SQLJsTransaction<TFullSchema, TSchema>) => T): T {
		const savepointName = `sp${this.nestedIndex + 1}`;
		const tx = new SQLJsTransaction('sync', this.dialect, this.session, this.schema, this.nestedIndex + 1);
		tx.run(sql.raw(`savepoint ${savepointName}`));
		try {
			const result = transaction(tx);
			tx.run(sql.raw(`release savepoint ${savepointName}`));
			return result;
		} catch (err) {
			tx.run(sql.raw(`rollback to savepoint ${savepointName}`));
			throw err;
		}
	}

Copyrun(placeholderValues?: Record<string, unknown>): void {
	const stmt = this.client.prepare(this.query.sql);
 
	const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
	this.logger.logQuery(this.query.sql, params);
	const result = stmt.run(params as BindParams);
 
	stmt.free();
 
	return result;
}
 
all(placeholderValues?: Record<string, unknown>): T['all'] {
	const stmt = this.client.prepare(this.query.sql);
 
	const { fields, joinsNotNullableMap, logger, query, customResultMapper } = this;
	if (!fields && !customResultMapper) {
		const params = fillPlaceholders(query.params, placeholderValues ?? {});
		logger.logQuery(query.sql, params);
		stmt.bind(params as BindParams);
		const rows: unknown[] = [];
		while (stmt.step()) {
			rows.push(stmt.getAsObject());
		}
 
		stmt.free();
 
		return rows;
	}
 
	const rows = this.values(placeholderValues) as unknown[][];
 
	if (customResultMapper) {
		return customResultMapper(rows, normalizeFieldValue) as T['all'];
	}
 
	return rows.map((row) => mapResultRow(fields!, row.map((v) => normalizeFieldValue(v)), joinsNotNullableMap));
}
 
get(placeholderValues?: Record<string, unknown>): T['get'] {
	const stmt = this.client.prepare(this.query.sql);
 
	const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
	this.logger.logQuery(this.query.sql, params);
 
	const { fields, joinsNotNullableMap, customResultMapper } = this;
	if (!fields && !customResultMapper) {
		const result = stmt.getAsObject(params as BindParams);
 
		stmt.free();
 
		return result;
	}
 
	const row = stmt.get(params as BindParams);
 
	stmt.free();
 
	if (!row || (row.length === 0 && fields!.length > 0)) {
		return undefined;
	}
 
	if (customResultMapper) {
		return customResultMapper([row], normalizeFieldValue) as T['get'];
	}
 
	return mapResultRow(fields!, row.map((v) => normalizeFieldValue(v)), joinsNotNullableMap);
}
 
values(placeholderValues?: Record<string, unknown>): T['values'] {
	const stmt = this.client.prepare(this.query.sql);
 
	const params = fillPlaceholders(this.query.params, placeholderValues ?? {});
	this.logger.logQuery(this.query.sql, params);
	stmt.bind(params as BindParams);
	const rows: unknown[] = [];
	while (stmt.step()) {
		rows.push(stmt.get());
	}
	stmt.free();
	return rows;
}

Copyfunction normalizeFieldValue(value: unknown) {
	if (value instanceof Uint8Array) { // eslint-disable-line no-instanceof/no-instanceof
		if (typeof Buffer !== 'undefined') {
			if (!(value instanceof Buffer)) { // eslint-disable-line no-instanceof/no-instanceof
				return Buffer.from(value);
			}
			return value;
		}
		if (typeof TextDecoder !== 'undefined') {
			return new TextDecoder().decode(value);
		}
		throw new Error('TextDecoder is not available. Please provide either Buffer or TextDecoder polyfill.');
	}
	return value;
}

Copyexport type AsyncRemoteCallback = (
	sql: string,
	params: any[],
	method: 'run' | 'all' | 'values' | 'get',
) => Promise<{ rows: any[] }>;
 
export type AsyncBatchRemoteCallback = (batch: {
	sql: string;
	params: any[];
	method: 'run' | 'all' | 'values' | 'get';
}[]) => Promise<{ rows: any[] }[]>;
 
export type RemoteCallback = AsyncRemoteCallback;
 
export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
	callback: RemoteCallback,
	config?: DrizzleConfig<TSchema>,
): SqliteRemoteDatabase<TSchema>;
export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
	callback: RemoteCallback,
	batchCallback?: AsyncBatchRemoteCallback,
	config?: DrizzleConfig<TSchema>,
): SqliteRemoteDatabase<TSchema>;
export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
	callback: RemoteCallback,
	batchCallback?: AsyncBatchRemoteCallback | DrizzleConfig<TSchema>,
	config?: DrizzleConfig<TSchema>,
): SqliteRemoteDatabase<TSchema> {

Copyconst dialect = new SQLiteAsyncDialect({ casing: config?.casing });
let logger;
let cache;
let _batchCallback: AsyncBatchRemoteCallback | undefined;
let _config: DrizzleConfig<TSchema> = {};
 
if (batchCallback) {
	if (typeof batchCallback === 'function') {
		_batchCallback = batchCallback as AsyncBatchRemoteCallback;
		_config = config ?? {};
	} else {
		_batchCallback = undefined;
		_config = batchCallback as DrizzleConfig<TSchema>;
	}
 
	if (_config.logger === true) {
		logger = new DefaultLogger();
	} else if (_config.logger !== false) {
		logger = _config.logger;
		cache = _config.cache;
	}
}
 
let schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined;
if (_config.schema) {
	const tablesConfig = extractTablesRelationalConfig(
		_config.schema,
		createTableRelationsHelpers,
	);
	schema = {
		fullSchema: _config.schema,
		schema: tablesConfig.tables,
		tableNamesMap: tablesConfig.tableNamesMap,
	};
}

Copyconst session = new SQLiteRemoteSession(callback, dialect, schema, _batchCallback, { logger, cache });
const db = new SqliteRemoteDatabase('async', dialect, session, schema) as SqliteRemoteDatabase<TSchema>;
(<any> db).$cache = cache;
if ((<any> db).$cache) {
	(<any> db).$cache['invalidate'] = cache?.onMutate;
}
return db;

Copyexport type AsyncRemoteCallback = (
	sql: string,
	params: any[],
	method: 'run' | 'all' | 'values' | 'get',
) => Promise<{ rows: any[] }>;
 
export type AsyncBatchRemoteCallback = (batch: {
	sql: string;
	params: any[];
	method: 'run' | 'all' | 'values' | 'get';
}[]) => Promise<{ rows: any[] }[]>;