Spaces:
Running
Running
| import { createPool as createCallbackPool, Pool as CallbackPool } from 'mysql2'; | |
| import { Pool, PoolConnection } from 'mysql2/promise'; | |
| import { | |
| AbstractSQLTeaQLClient, | |
| assertSafeIdentifier, | |
| ColumnSchema, | |
| EntitySchema, | |
| LogicalColumnType, | |
| SqlQueryResult, | |
| SqlSession, | |
| standardAggregateFunction, | |
| ensureOptimisticIdFloor, | |
| TeaQLSqlDriver, | |
| canonicalRelationIndexes, | |
| } from './core'; | |
| class MySQLSession implements SqlSession { | |
| constructor(private readonly client: Pool | PoolConnection) {} | |
| async query(sql: string, values: any[] = []): Promise<SqlQueryResult> { | |
| const [raw]: any = await (this.client as any).query(sql, values); | |
| if (Array.isArray(raw)) return { rows: raw, rowCount: raw.length }; | |
| return { rows: [], rowCount: Number(raw?.affectedRows || 0) }; | |
| } | |
| } | |
| export class MySQLDriver implements TeaQLSqlDriver { | |
| readonly databaseKind = 'mysql' as const; | |
| private readonly callbackPool: CallbackPool; | |
| private readonly pool: Pool; | |
| constructor(connectionString: string) { | |
| if (!connectionString) throw new Error('connectionString is required'); | |
| this.callbackPool = createCallbackPool({ uri: connectionString, timezone: 'Z' }); | |
| this.pool = this.callbackPool.promise(); | |
| } | |
| identifier(value: string): string { | |
| return `\`${assertSafeIdentifier(value)}\``; | |
| } | |
| placeholder(_index: number): string { | |
| return '?'; | |
| } | |
| encode(value: any, column?: ColumnSchema): any { | |
| if (value === null || value === undefined) return value; | |
| if (column?.logicalType === 'json' && typeof value !== 'string') { | |
| return JSON.stringify(value); | |
| } | |
| if (column?.logicalType === 'boolean') return value ? 1 : 0; | |
| return value; | |
| } | |
| contains(columnSql: string, placeholder: string): string { | |
| return `CAST(${columnSql} AS CHAR) LIKE CONCAT('%', ${placeholder}, '%')`; | |
| } | |
| aggregateFunction(name: string): string { | |
| const functions: Record<string, string> = { | |
| stddev: 'STDDEV_SAMP', | |
| stddevpop: 'STDDEV_POP', | |
| varsamp: 'VAR_SAMP', | |
| varpop: 'VAR_POP', | |
| bitand: 'BIT_AND', | |
| bitor: 'BIT_OR', | |
| bitxor: 'BIT_XOR', | |
| }; | |
| return functions[String(name).toLowerCase()] | |
| || standardAggregateFunction(name); | |
| } | |
| private sqlType(type: LogicalColumnType): string { | |
| const types: Record<LogicalColumnType, string> = { | |
| boolean: 'BOOLEAN', | |
| double: 'DOUBLE', | |
| decimal: 'DECIMAL(65, 30)', | |
| date: 'DATE', | |
| datetime: 'DATETIME(6)', | |
| json: 'JSON', | |
| integer: 'BIGINT', | |
| text: 'TEXT', | |
| }; | |
| return types[type]; | |
| } | |
| async ensureSchema(schemas: Record<string, EntitySchema>): Promise<void> { | |
| for (const schema of Object.values(schemas)) { | |
| const table = this.identifier(schema.table); | |
| await this.query( | |
| `CREATE TABLE IF NOT EXISTS ${table} (` + | |
| '`id` BIGINT PRIMARY KEY, `version` BIGINT NOT NULL)', | |
| ); | |
| for (const [field, column] of Object.entries(schema.columns)) { | |
| if (field === 'id' || field === 'version') continue; | |
| const existing = await this.query( | |
| 'SELECT 1 FROM information_schema.columns ' + | |
| 'WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?', | |
| [schema.table, column.columnName], | |
| ); | |
| if (!existing.rowCount) { | |
| await this.query( | |
| `ALTER TABLE ${table} ADD COLUMN ` + | |
| `${this.identifier(column.columnName)} ${this.sqlType(column.logicalType)}` + | |
| `${column.nullable === false ? ' NOT NULL' : ''}`, | |
| ); | |
| } | |
| } | |
| } | |
| for (const index of canonicalRelationIndexes(schemas)) { | |
| const existing = await this.query( | |
| 'SELECT 1 FROM information_schema.statistics ' + | |
| 'WHERE table_schema = DATABASE() AND table_name = ? AND index_name = ?', | |
| [index.table, index.name], | |
| ); | |
| if (!existing.rowCount) { | |
| await this.query( | |
| `CREATE INDEX ${this.identifier(index.name)} ON ` + | |
| `${this.identifier(index.table)} (` + | |
| `${this.identifier(index.foreignColumn)}, ${this.identifier(index.idColumn)} DESC)`, | |
| ); | |
| } | |
| } | |
| await this.query( | |
| 'CREATE TABLE IF NOT EXISTS teaql_id_space (' + | |
| 'type_name VARCHAR(255) PRIMARY KEY, current_level BIGINT NOT NULL)', | |
| ); | |
| } | |
| async transaction<T>(work: (session: SqlSession) => Promise<T>): Promise<T> { | |
| const connection = await this.pool.getConnection(); | |
| try { | |
| await connection.beginTransaction(); | |
| const result = await work(new MySQLSession(connection)); | |
| await connection.commit(); | |
| return result; | |
| } catch (error) { | |
| await connection.rollback(); | |
| throw error; | |
| } finally { | |
| connection.release(); | |
| } | |
| } | |
| async nextId(session: SqlSession, entity: string): Promise<string> { | |
| for (let attempt = 1; attempt <= 100; attempt += 1) { | |
| const result = await session.query( | |
| 'SELECT current_level AS id FROM teaql_id_space WHERE type_name = ?', [entity]); | |
| if (!result.rowCount) { | |
| try { | |
| const inserted = await session.query( | |
| 'INSERT INTO teaql_id_space(type_name, current_level) VALUES (?, 1)', [entity]); | |
| if (inserted.rowCount === 1) return '1'; | |
| } catch (error) { | |
| const winner = await session.query( | |
| 'SELECT current_level FROM teaql_id_space WHERE type_name = ?', [entity]); | |
| if (!winner.rowCount) throw error; | |
| } | |
| continue; | |
| } | |
| const previous = Number(result.rows[0].id); | |
| const next = previous + 1; | |
| const updated = await session.query( | |
| 'UPDATE teaql_id_space SET current_level = ? ' + | |
| 'WHERE type_name = ? AND current_level = ?', [next, entity, previous]); | |
| if (updated.rowCount === 1) return String(next); | |
| if (updated.rowCount !== 0) throw new Error( | |
| `ID space update for ${entity} changed ${updated.rowCount} rows`); | |
| } | |
| throw new Error(`Unable to allocate ID for ${entity} after 100 optimistic-lock attempts`); | |
| } | |
| ensureIdFloor(session: SqlSession, entity: string, floor: string): Promise<void> { | |
| return ensureOptimisticIdFloor(session, index => this.placeholder(index), entity, floor); | |
| } | |
| query(sql: string, values: any[] = []): Promise<SqlQueryResult> { | |
| return new MySQLSession(this.pool).query(sql, values); | |
| } | |
| async *stream(sql: string, values: any[] = []): AsyncIterable<any> { | |
| const connection = await new Promise<any>((resolve, reject) => { | |
| this.callbackPool.getConnection((error, value) => error ? reject(error) : resolve(value)); | |
| }); | |
| const readable = connection.query(sql, values).stream(); | |
| try { | |
| for await (const row of readable) yield row; | |
| } finally { | |
| readable.destroy(); | |
| connection.release(); | |
| } | |
| } | |
| async close(): Promise<void> { | |
| await this.callbackPool.promise().end(); | |
| } | |
| } | |
| export class MySQLTeaQLClient extends AbstractSQLTeaQLClient { | |
| constructor(connectionString: string, schemas: Record<string, EntitySchema>) { | |
| super(new MySQLDriver(connectionString), schemas); | |
| this.attachStandardTelemetryStartup(startStandardNodeTelemetry()); | |
| } | |
| } | |
| import { startStandardNodeTelemetry } from '../telemetry/standard-node'; | |