Add TiDB Serverless driver

.nvmrc

@@ -1 +1 @@
-18
+18.18

changelogs/drizzle-orm/0.31.1.md

@@ -0,0 +1,9 @@
+- 🎉 Added support for TiDB Cloud Serverless driver:
+  ```ts
+  import { connect } from '@tidbcloud/serverless';
+  import { drizzle } from 'drizzle-orm/tidb-serverless';
+
+  const client = connect({ url: '...' });
+  const db = drizzle(client);
+  await db.select().from(...);
+  ```

drizzle-orm/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "drizzle-orm",
-	"version": "0.31.0",
+	"version": "0.31.1",
 	"description": "Drizzle ORM package for SQL databases",
 	"type": "module",
 	"scripts": {
@@ -67,7 +67,8 @@
 		"postgres": ">=3",
 		"react": ">=18",
 		"sql.js": ">=1",
-		"sqlite3": ">=5"
+		"sqlite3": ">=5",
+		"@tidbcloud/serverless": "*"
 	},
 	"peerDependenciesMeta": {
 		"mysql2": {
@@ -144,6 +145,9 @@
 		},
 		"@electric-sql/pglite": {
 			"optional": true
+		},
+		"@tidbcloud/serverless": {
+			"optional": true
 		}
 	},
 	"devDependencies": {
@@ -156,6 +160,7 @@
 		"@opentelemetry/api": "^1.4.1",
 		"@originjs/vite-plugin-commonjs": "^1.0.3",
 		"@planetscale/database": "^1.16.0",
+		"@tidbcloud/serverless": "^0.1.1",
 		"@types/better-sqlite3": "^7.6.4",
 		"@types/node": "^20.2.5",
 		"@types/pg": "^8.10.1",

drizzle-orm/src/tidb-serverless/driver.ts

@@ -0,0 +1,51 @@
+import type { Connection } from '@tidbcloud/serverless';
+import type { Logger } from '~/logger.ts';
+import { DefaultLogger } from '~/logger.ts';
+import { MySqlDatabase } from '~/mysql-core/db.ts';
+import { MySqlDialect } from '~/mysql-core/dialect.ts';
+import {
+	createTableRelationsHelpers,
+	extractTablesRelationalConfig,
+	type RelationalSchemaConfig,
+	type TablesRelationalConfig,
+} from '~/relations.ts';
+import type { DrizzleConfig } from '~/utils.ts';
+import type { TiDBServerlessPreparedQueryHKT, TiDBServerlessQueryResultHKT } from './session.ts';
+import { TiDBServerlessSession } from './session.ts';
+
+export interface TiDBServerlessSDriverOptions {
+	logger?: Logger;
+}
+
+export type TiDBServerlessDatabase<
+	TSchema extends Record<string, unknown> = Record<string, never>,
+> = MySqlDatabase<TiDBServerlessQueryResultHKT, TiDBServerlessPreparedQueryHKT, TSchema>;
+
+export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
+	client: Connection,
+	config: DrizzleConfig<TSchema> = {},
+): TiDBServerlessDatabase<TSchema> {
+	const dialect = new MySqlDialect();
+	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 TiDBServerlessSession(client, dialect, undefined, schema, { logger });
+	return new MySqlDatabase(dialect, session, schema, 'default') as TiDBServerlessDatabase<TSchema>;
+}

drizzle-orm/src/tidb-serverless/index.ts

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

drizzle-orm/src/tidb-serverless/migrator.ts

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

drizzle-orm/src/tidb-serverless/session.ts

@@ -0,0 +1,171 @@
+import type { Connection, ExecuteOptions, FullResult, Tx } from '@tidbcloud/serverless';
+
+import { entityKind } from '~/entity.ts';
+import type { Logger } from '~/logger.ts';
+import { NoopLogger } from '~/logger.ts';
+import type { MySqlDialect } from '~/mysql-core/dialect.ts';
+import type { SelectedFieldsOrdered } from '~/mysql-core/query-builders/select.types.ts';
+import {
+	MySqlSession,
+	MySqlTransaction,
+	PreparedQuery,
+	type PreparedQueryConfig,
+	type PreparedQueryHKT,
+	type QueryResultHKT,
+} from '~/mysql-core/session.ts';
+import type { RelationalSchemaConfig, TablesRelationalConfig } from '~/relations.ts';
+import { fillPlaceholders, type Query, type SQL, sql } from '~/sql/sql.ts';
+import { type Assume, mapResultRow } from '~/utils.ts';
+
+const executeRawConfig = { fullResult: true } satisfies ExecuteOptions;
+const queryConfig = { arrayMode: true } satisfies ExecuteOptions;
+
+export class TiDBServerlessPreparedQuery<T extends PreparedQueryConfig> extends PreparedQuery<T> {
+	static readonly [entityKind]: string = 'TiDBPreparedQuery';
+
+	constructor(
+		private client: Tx | Connection,
+		private queryString: string,
+		private params: unknown[],
+		private logger: Logger,
+		private fields: SelectedFieldsOrdered | undefined,
+		private customResultMapper?: (rows: unknown[][]) => T['execute'],
+	) {
+		super();
+	}
+
+	async execute(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['execute']> {
+		const params = fillPlaceholders(this.params, placeholderValues);
+
+		this.logger.logQuery(this.queryString, params);
+
+		const { fields, client, queryString, joinsNotNullableMap, customResultMapper } = this;
+		if (!fields && !customResultMapper) {
+			return client.execute(queryString, params, executeRawConfig);
+		}
+
+		const rows = await client.execute(queryString, params, queryConfig) as unknown[][];
+
+		if (customResultMapper) {
+			return customResultMapper(rows);
+		}
+
+		return rows.map((row) => mapResultRow<T['execute']>(fields!, row, joinsNotNullableMap));
+	}
+
+	override iterator(_placeholderValues?: Record<string, unknown>): AsyncGenerator<T['iterator']> {
+		throw new Error('Streaming is not supported by the TiDB Cloud Serverless driver');
+	}
+}
+
+export interface TiDBServerlessSessionOptions {
+	logger?: Logger;
+}
+
+export class TiDBServerlessSession<
+	TFullSchema extends Record<string, unknown>,
+	TSchema extends TablesRelationalConfig,
+> extends MySqlSession<TiDBServerlessQueryResultHKT, TiDBServerlessPreparedQueryHKT, TFullSchema, TSchema> {
+	static readonly [entityKind]: string = 'TiDBServerlessSession';
+
+	private logger: Logger;
+	private client: Tx | Connection;
+
+	constructor(
+		private baseClient: Connection,
+		dialect: MySqlDialect,
+		tx: Tx | undefined,
+		private schema: RelationalSchemaConfig<TSchema> | undefined,
+		private options: TiDBServerlessSessionOptions = {},
+	) {
+		super(dialect);
+		this.client = tx ?? baseClient;
+		this.logger = options.logger ?? new NoopLogger();
+	}
+
+	prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
+		query: Query,
+		fields: SelectedFieldsOrdered | undefined,
+		customResultMapper?: (rows: unknown[][]) => T['execute'],
+	): PreparedQuery<T> {
+		return new TiDBServerlessPreparedQuery(
+			this.client,
+			query.sql,
+			query.params,
+			this.logger,
+			fields,
+			customResultMapper,
+		);
+	}
+
+	override all<T = unknown>(query: SQL): Promise<T[]> {
+		const querySql = this.dialect.sqlToQuery(query);
+		this.logger.logQuery(querySql.sql, querySql.params);
+		return this.client.execute(querySql.sql, querySql.params) as Promise<T[]>;
+	}
+
+	override async transaction<T>(
+		transaction: (tx: TiDBServerlessTransaction<TFullSchema, TSchema>) => Promise<T>,
+	): Promise<T> {
+		const nativeTx = await this.baseClient.begin();
+		try {
+			const session = new TiDBServerlessSession(this.baseClient, this.dialect, nativeTx, this.schema, this.options);
+			const tx = new TiDBServerlessTransaction<TFullSchema, TSchema>(
+				this.dialect,
+				session as MySqlSession<any, any, any, any>,
+				this.schema,
+			);
+			const result = await transaction(tx);
+			await nativeTx.commit();
+			return result;
+		} catch (err) {
+			await nativeTx.rollback();
+			throw err;
+		}
+	}
+}
+
+export class TiDBServerlessTransaction<
+	TFullSchema extends Record<string, unknown>,
+	TSchema extends TablesRelationalConfig,
+> extends MySqlTransaction<TiDBServerlessQueryResultHKT, TiDBServerlessPreparedQueryHKT, TFullSchema, TSchema> {
+	static readonly [entityKind]: string = 'TiDBServerlessTransaction';
+
+	constructor(
+		dialect: MySqlDialect,
+		session: MySqlSession,
+		schema: RelationalSchemaConfig<TSchema> | undefined,
+		nestedIndex = 0,
+	) {
+		super(dialect, session, schema, nestedIndex, 'default');
+	}
+
+	override async transaction<T>(
+		transaction: (tx: TiDBServerlessTransaction<TFullSchema, TSchema>) => Promise<T>,
+	): Promise<T> {
+		const savepointName = `sp${this.nestedIndex + 1}`;
+		const tx = new TiDBServerlessTransaction<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 TiDBServerlessQueryResultHKT extends QueryResultHKT {
+	type: FullResult;
+}
+
+export interface TiDBServerlessPreparedQueryHKT extends PreparedQueryHKT {
+	type: TiDBServerlessPreparedQuery<Assume<this['config'], PreparedQueryConfig>>;
+}

integration-tests/.env.example

@@ -1,6 +1,7 @@
 PG_CONNECTION_STRING="postgres://postgres:postgres@localhost:55432/postgres"
 MYSQL_CONNECTION_STRING="mysql://root:mysql@127.0.0.1:33306/drizzle"
 PLANETSCALE_CONNECTION_STRING=
+TIDB_CONNECTION_STRING=
 NEON_CONNECTION_STRING=
 LIBSQL_URL="file:local.db"
 LIBSQL_AUTH_TOKEN="ey..." # For Turso only

integration-tests/package.json

@@ -28,6 +28,7 @@
 			"!tests/sqlite-proxy-batch.test.ts",
 			"!tests/neon-http-batch.test.ts",
 			"!tests/neon-http.test.ts",
+			"!tests/tidb-serverless.test.ts",
 			"!tests/replicas/**/*",
 			"!tests/imports/**/*",
 			"!tests/extensions/**/*"
@@ -70,6 +71,7 @@
 		"@miniflare/d1": "^2.14.2",
 		"@miniflare/shared": "^2.14.2",
 		"@planetscale/database": "^1.16.0",
+		"@tidbcloud/serverless": "^0.1.1",
 		"@typescript/analyze-trace": "^0.10.0",
 		"@vercel/postgres": "^0.3.0",
 		"@xata.io/client": "^0.29.3",