Skip to content

Commit 26a7171

Browse files
dankochetovdankochetov
committed
Add TiDB Serverless driver
1 parent 6df4b83 commit 26a7171

File tree

11 files changed

+13159
-7778
lines changed

11 files changed

+13159
-7778
lines changed

‎.nvmrc

+1-1
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
18
1+
18.18

‎changelogs/drizzle-orm/0.31.1.md

+9
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
- 🎉 Added support for TiDB Cloud Serverless driver:
2+
```ts
3+
import { connect } from '@tidbcloud/serverless';
4+
import { drizzle } from 'drizzle-orm/tidb-serverless';
5+
6+
const client = connect({ url: '...' });
7+
const db = drizzle(client);
8+
await db.select().from(...);
9+
```

‎drizzle-orm/package.json

+7-2
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "drizzle-orm",
3-
"version": "0.31.0",
3+
"version": "0.31.1",
44
"description": "Drizzle ORM package for SQL databases",
55
"type": "module",
66
"scripts": {
@@ -67,7 +67,8 @@
6767
"postgres": ">=3",
6868
"react": ">=18",
6969
"sql.js": ">=1",
70-
"sqlite3": ">=5"
70+
"sqlite3": ">=5",
71+
"@tidbcloud/serverless": "*"
7172
},
7273
"peerDependenciesMeta": {
7374
"mysql2": {
@@ -144,6 +145,9 @@
144145
},
145146
"@electric-sql/pglite": {
146147
"optional": true
148+
},
149+
"@tidbcloud/serverless": {
150+
"optional": true
147151
}
148152
},
149153
"devDependencies": {
@@ -156,6 +160,7 @@
156160
"@opentelemetry/api": "^1.4.1",
157161
"@originjs/vite-plugin-commonjs": "^1.0.3",
158162
"@planetscale/database": "^1.16.0",
163+
"@tidbcloud/serverless": "^0.1.1",
159164
"@types/better-sqlite3": "^7.6.4",
160165
"@types/node": "^20.2.5",
161166
"@types/pg": "^8.10.1",

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

+51
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import type { Connection } from '@tidbcloud/serverless';
2+
import type { Logger } from '~/logger.ts';
3+
import { DefaultLogger } from '~/logger.ts';
4+
import { MySqlDatabase } from '~/mysql-core/db.ts';
5+
import { MySqlDialect } from '~/mysql-core/dialect.ts';
6+
import {
7+
createTableRelationsHelpers,
8+
extractTablesRelationalConfig,
9+
type RelationalSchemaConfig,
10+
type TablesRelationalConfig,
11+
} from '~/relations.ts';
12+
import type { DrizzleConfig } from '~/utils.ts';
13+
import type { TiDBServerlessPreparedQueryHKT, TiDBServerlessQueryResultHKT } from './session.ts';
14+
import { TiDBServerlessSession } from './session.ts';
15+
16+
export interface TiDBServerlessSDriverOptions {
17+
logger?: Logger;
18+
}
19+
20+
export type TiDBServerlessDatabase<
21+
TSchema extends Record<string, unknown> = Record<string, never>,
22+
> = MySqlDatabase<TiDBServerlessQueryResultHKT, TiDBServerlessPreparedQueryHKT, TSchema>;
23+
24+
export function drizzle<TSchema extends Record<string, unknown> = Record<string, never>>(
25+
client: Connection,
26+
config: DrizzleConfig<TSchema> = {},
27+
): TiDBServerlessDatabase<TSchema> {
28+
const dialect = new MySqlDialect();
29+
let logger;
30+
if (config.logger === true) {
31+
logger = new DefaultLogger();
32+
} else if (config.logger !== false) {
33+
logger = config.logger;
34+
}
35+
36+
let schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined;
37+
if (config.schema) {
38+
const tablesConfig = extractTablesRelationalConfig(
39+
config.schema,
40+
createTableRelationsHelpers,
41+
);
42+
schema = {
43+
fullSchema: config.schema,
44+
schema: tablesConfig.tables,
45+
tableNamesMap: tablesConfig.tableNamesMap,
46+
};
47+
}
48+
49+
const session = new TiDBServerlessSession(client, dialect, undefined, schema, { logger });
50+
return new MySqlDatabase(dialect, session, schema, 'default') as TiDBServerlessDatabase<TSchema>;
51+
}

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

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
export * from './driver.ts';
2+
export * from './session.ts';

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

+11
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
import type { MigrationConfig } from '~/migrator.ts';
2+
import { readMigrationFiles } from '~/migrator.ts';
3+
import type { TiDBServerlessDatabase } from './driver.ts';
4+
5+
export async function migrate<TSchema extends Record<string, unknown>>(
6+
db: TiDBServerlessDatabase<TSchema>,
7+
config: MigrationConfig,
8+
) {
9+
const migrations = readMigrationFiles(config);
10+
await db.dialect.migrate(migrations, db.session, config);
11+
}

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

+171
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
import type { Connection, ExecuteOptions, FullResult, Tx } from '@tidbcloud/serverless';
2+
3+
import { entityKind } from '~/entity.ts';
4+
import type { Logger } from '~/logger.ts';
5+
import { NoopLogger } from '~/logger.ts';
6+
import type { MySqlDialect } from '~/mysql-core/dialect.ts';
7+
import type { SelectedFieldsOrdered } from '~/mysql-core/query-builders/select.types.ts';
8+
import {
9+
MySqlSession,
10+
MySqlTransaction,
11+
PreparedQuery,
12+
type PreparedQueryConfig,
13+
type PreparedQueryHKT,
14+
type QueryResultHKT,
15+
} from '~/mysql-core/session.ts';
16+
import type { RelationalSchemaConfig, TablesRelationalConfig } from '~/relations.ts';
17+
import { fillPlaceholders, type Query, type SQL, sql } from '~/sql/sql.ts';
18+
import { type Assume, mapResultRow } from '~/utils.ts';
19+
20+
const executeRawConfig = { fullResult: true } satisfies ExecuteOptions;
21+
const queryConfig = { arrayMode: true } satisfies ExecuteOptions;
22+
23+
export class TiDBServerlessPreparedQuery<T extends PreparedQueryConfig> extends PreparedQuery<T> {
24+
static readonly [entityKind]: string = 'TiDBPreparedQuery';
25+
26+
constructor(
27+
private client: Tx | Connection,
28+
private queryString: string,
29+
private params: unknown[],
30+
private logger: Logger,
31+
private fields: SelectedFieldsOrdered | undefined,
32+
private customResultMapper?: (rows: unknown[][]) => T['execute'],
33+
) {
34+
super();
35+
}
36+
37+
async execute(placeholderValues: Record<string, unknown> | undefined = {}): Promise<T['execute']> {
38+
const params = fillPlaceholders(this.params, placeholderValues);
39+
40+
this.logger.logQuery(this.queryString, params);
41+
42+
const { fields, client, queryString, joinsNotNullableMap, customResultMapper } = this;
43+
if (!fields && !customResultMapper) {
44+
return client.execute(queryString, params, executeRawConfig);
45+
}
46+
47+
const rows = await client.execute(queryString, params, queryConfig) as unknown[][];
48+
49+
if (customResultMapper) {
50+
return customResultMapper(rows);
51+
}
52+
53+
return rows.map((row) => mapResultRow<T['execute']>(fields!, row, joinsNotNullableMap));
54+
}
55+
56+
override iterator(_placeholderValues?: Record<string, unknown>): AsyncGenerator<T['iterator']> {
57+
throw new Error('Streaming is not supported by the TiDB Cloud Serverless driver');
58+
}
59+
}
60+
61+
export interface TiDBServerlessSessionOptions {
62+
logger?: Logger;
63+
}
64+
65+
export class TiDBServerlessSession<
66+
TFullSchema extends Record<string, unknown>,
67+
TSchema extends TablesRelationalConfig,
68+
> extends MySqlSession<TiDBServerlessQueryResultHKT, TiDBServerlessPreparedQueryHKT, TFullSchema, TSchema> {
69+
static readonly [entityKind]: string = 'TiDBServerlessSession';
70+
71+
private logger: Logger;
72+
private client: Tx | Connection;
73+
74+
constructor(
75+
private baseClient: Connection,
76+
dialect: MySqlDialect,
77+
tx: Tx | undefined,
78+
private schema: RelationalSchemaConfig<TSchema> | undefined,
79+
private options: TiDBServerlessSessionOptions = {},
80+
) {
81+
super(dialect);
82+
this.client = tx ?? baseClient;
83+
this.logger = options.logger ?? new NoopLogger();
84+
}
85+
86+
prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
87+
query: Query,
88+
fields: SelectedFieldsOrdered | undefined,
89+
customResultMapper?: (rows: unknown[][]) => T['execute'],
90+
): PreparedQuery<T> {
91+
return new TiDBServerlessPreparedQuery(
92+
this.client,
93+
query.sql,
94+
query.params,
95+
this.logger,
96+
fields,
97+
customResultMapper,
98+
);
99+
}
100+
101+
override all<T = unknown>(query: SQL): Promise<T[]> {
102+
const querySql = this.dialect.sqlToQuery(query);
103+
this.logger.logQuery(querySql.sql, querySql.params);
104+
return this.client.execute(querySql.sql, querySql.params) as Promise<T[]>;
105+
}
106+
107+
override async transaction<T>(
108+
transaction: (tx: TiDBServerlessTransaction<TFullSchema, TSchema>) => Promise<T>,
109+
): Promise<T> {
110+
const nativeTx = await this.baseClient.begin();
111+
try {
112+
const session = new TiDBServerlessSession(this.baseClient, this.dialect, nativeTx, this.schema, this.options);
113+
const tx = new TiDBServerlessTransaction<TFullSchema, TSchema>(
114+
this.dialect,
115+
session as MySqlSession<any, any, any, any>,
116+
this.schema,
117+
);
118+
const result = await transaction(tx);
119+
await nativeTx.commit();
120+
return result;
121+
} catch (err) {
122+
await nativeTx.rollback();
123+
throw err;
124+
}
125+
}
126+
}
127+
128+
export class TiDBServerlessTransaction<
129+
TFullSchema extends Record<string, unknown>,
130+
TSchema extends TablesRelationalConfig,
131+
> extends MySqlTransaction<TiDBServerlessQueryResultHKT, TiDBServerlessPreparedQueryHKT, TFullSchema, TSchema> {
132+
static readonly [entityKind]: string = 'TiDBServerlessTransaction';
133+
134+
constructor(
135+
dialect: MySqlDialect,
136+
session: MySqlSession,
137+
schema: RelationalSchemaConfig<TSchema> | undefined,
138+
nestedIndex = 0,
139+
) {
140+
super(dialect, session, schema, nestedIndex, 'default');
141+
}
142+
143+
override async transaction<T>(
144+
transaction: (tx: TiDBServerlessTransaction<TFullSchema, TSchema>) => Promise<T>,
145+
): Promise<T> {
146+
const savepointName = `sp${this.nestedIndex + 1}`;
147+
const tx = new TiDBServerlessTransaction<TFullSchema, TSchema>(
148+
this.dialect,
149+
this.session,
150+
this.schema,
151+
this.nestedIndex + 1,
152+
);
153+
await tx.execute(sql.raw(`savepoint ${savepointName}`));
154+
try {
155+
const result = await transaction(tx);
156+
await tx.execute(sql.raw(`release savepoint ${savepointName}`));
157+
return result;
158+
} catch (err) {
159+
await tx.execute(sql.raw(`rollback to savepoint ${savepointName}`));
160+
throw err;
161+
}
162+
}
163+
}
164+
165+
export interface TiDBServerlessQueryResultHKT extends QueryResultHKT {
166+
type: FullResult;
167+
}
168+
169+
export interface TiDBServerlessPreparedQueryHKT extends PreparedQueryHKT {
170+
type: TiDBServerlessPreparedQuery<Assume<this['config'], PreparedQueryConfig>>;
171+
}

‎integration-tests/.env.example

+1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
PG_CONNECTION_STRING="postgres://postgres:postgres@localhost:55432/postgres"
22
MYSQL_CONNECTION_STRING="mysql://root:[email protected]:33306/drizzle"
33
PLANETSCALE_CONNECTION_STRING=
4+
TIDB_CONNECTION_STRING=
45
NEON_CONNECTION_STRING=
56
LIBSQL_URL="file:local.db"
67
LIBSQL_AUTH_TOKEN="ey..." # For Turso only

‎integration-tests/package.json

+2
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
"!tests/sqlite-proxy-batch.test.ts",
2929
"!tests/neon-http-batch.test.ts",
3030
"!tests/neon-http.test.ts",
31+
"!tests/tidb-serverless.test.ts",
3132
"!tests/replicas/**/*",
3233
"!tests/imports/**/*",
3334
"!tests/extensions/**/*"
@@ -70,6 +71,7 @@
7071
"@miniflare/d1": "^2.14.2",
7172
"@miniflare/shared": "^2.14.2",
7273
"@planetscale/database": "^1.16.0",
74+
"@tidbcloud/serverless": "^0.1.1",
7375
"@typescript/analyze-trace": "^0.10.0",
7476
"@vercel/postgres": "^0.3.0",
7577
"@xata.io/client": "^0.29.3",

0 commit comments

Comments
 (0)