TiDBServerlessSession Class — drizzle-orm Architecture

Source Code

drizzle-orm/src/tidb-serverless/session.ts lines 109–192

export class TiDBServerlessSession<
	TFullSchema extends Record<string, unknown>,
	TSchema extends TablesRelationalConfig,
> extends MySqlSession<TiDBServerlessQueryResultHKT, TiDBServerlessPreparedQueryHKT, TFullSchema, TSchema> {
	static override readonly [entityKind]: string = 'TiDBServerlessSession';

	private logger: Logger;
	private client: Tx | Connection;
	private cache: Cache;

	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();
		this.cache = options.cache ?? new NoopCache();
	}

	prepareQuery<T extends MySqlPreparedQueryConfig = MySqlPreparedQueryConfig>(
		query: Query,
		fields: SelectedFieldsOrdered | undefined,
		customResultMapper?: (rows: unknown[][]) => T['execute'],
		generatedIds?: Record<string, unknown>[],
		returningIds?: SelectedFieldsOrdered,
		queryMetadata?: {
			type: 'select' | 'update' | 'delete' | 'insert';
			tables: string[];
		},
		cacheConfig?: WithCacheConfig,
	): MySqlPreparedQuery<T> {
		return new TiDBServerlessPreparedQuery(
			this.client,
			query.sql,
			query.params,
			this.logger,
			this.cache,
			queryMetadata,
			cacheConfig,
			fields,
			customResultMapper,
			generatedIds,
			returningIds,
		);
	}

	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 count(sql: SQL): Promise<number> {
		const res = await this.execute<{ rows: [{ count: string }] }>(sql);

		return Number(
			res['rows'][0]['count'],
		);
	}

	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;