...

Author: igalklebanov

src/dialect/mysql/mysql-dialect-config.ts

@@ -6,6 +6,13 @@ import { DatabaseConnection } from '../../driver/database-connection.js'
  * https://github.com/sidorares/node-mysql2#using-connection-pools
  */
 export interface MysqlDialectConfig {
+  /**
+   * The `mysql2` `createConnection` function or similar.
+   *
+   * TODO: ...
+   */
+  createConnection?: (opts: any) => MysqlConnection | Promise<MysqlConnection>
+
   /**
    * A mysql2 Pool instance or a function that returns one.
    *
@@ -41,7 +48,9 @@ export interface MysqlPool {
   end(callback: (error: unknown) => void): void
 }
 
-export interface MysqlPoolConnection {
+export interface MysqlConnection {
+  config: unknown
+  connect(callback?: (error: unknown) => void): void
   destroy(): void
   query(
     sql: string,
@@ -52,10 +61,15 @@ export interface MysqlPoolConnection {
     parameters: ReadonlyArray<unknown>,
     callback: (error: unknown, result: MysqlQueryResult) => void,
   ): void
-  release(): void
   threadId: number
 }
 
+export type MysqlConectionConstructor = new (opts?: object) => MysqlConnection
+
+export interface MysqlPoolConnection extends MysqlConnection {
+  release(): void
+}
+
 export interface MysqlStreamOptions {
   highWaterMark?: number
   objectMode?: boolean

src/dialect/mysql/mysql-driver.ts

@@ -42,7 +42,10 @@ export class MysqlDriver implements Driver {
     let connection = this.#connections.get(rawConnection)
 
     if (!connection) {
-      connection = new MysqlConnection(rawConnection)
+      connection = new MysqlConnection(
+        rawConnection,
+        this.#config.createConnection,
+      )
       this.#connections.set(rawConnection, connection)
 
       // The driver must take care of calling `onCreateConnection` when a new
@@ -165,9 +168,14 @@ function isOkPacket(obj: unknown): obj is MysqlOkPacket {
 }
 
 class MysqlConnection implements DatabaseConnection {
+  readonly #createConnection: MysqlDialectConfig['createConnection']
   readonly #rawConnection: MysqlPoolConnection
 
-  constructor(rawConnection: MysqlPoolConnection) {
+  constructor(
+    rawConnection: MysqlPoolConnection,
+    createConnection: MysqlDialectConfig['createConnection'],
+  ) {
+    this.#createConnection = createConnection
     this.#rawConnection = rawConnection
   }
 
@@ -183,14 +191,38 @@ class MysqlConnection implements DatabaseConnection {
       // noop
     }
 
+    const { config, threadId } = this.#rawConnection
+
+    // this kills the query and the connection database-side.
+    // we're not using `kill query <connection_id>` here because it doesn't
+    // guarantee that the query is killed immediately. we saw that in tests,
+    // the query can still run for a while after - including registering writes.
+    const cancelQuery = `kill connection ${threadId}`
+
+    if (this.#createConnection && config) {
+      const controlConnection = await this.#createConnection({ ...config })
+
+      return await new Promise((resolve, reject) => {
+        controlConnection.connect((connectError) => {
+          if (connectError) {
+            return reject(connectError)
+          }
+
+          controlConnection.query(cancelQuery, [], (error) => {
+            controlConnection.destroy()
+
+            if (error) {
+              return reject(error)
+            }
+
+            resolve()
+          })
+        })
+      })
+    }
+
     await controlConnectionProvider(async (controlConnection) => {
-      // this kills the query and the connection database-side.
-      // we're not using `kill query <connection_id>` here because it doesn't
-      // guarantee that the query is killed immediately. we saw that in tests,
-      // the query can still run for a while after - including registering writes.
-      await controlConnection.executeQuery(
-        CompiledQuery.raw('kill connection ?', [this.#rawConnection.threadId]),
-      )
+      await controlConnection.executeQuery(CompiledQuery.raw(cancelQuery, []))
     })
   }
 
@@ -231,7 +263,6 @@ class MysqlConnection implements DatabaseConnection {
   async *streamQuery<O>(
     compiledQuery: CompiledQuery,
     _chunkSize: number,
-    options?: QueryOptions,
   ): AsyncIterableIterator<QueryResult<O>> {
     const stream = this.#rawConnection
       .query(compiledQuery.sql, compiledQuery.parameters)

test/node/src/test-setup.ts

@@ -3,7 +3,7 @@ import chaiAsPromised from 'chai-as-promised'
 import * as chaiSubset from 'chai-subset'
 import * as Cursor from 'pg-cursor'
 import { Pool, PoolConfig } from 'pg'
-import { createPool } from 'mysql2'
+import { createConnection, createPool } from 'mysql2'
 import * as Database from 'better-sqlite3'
 import * as Tarn from 'tarn'
 import * as Tedious from 'tedious'
@@ -182,6 +182,7 @@ export const DB_CONFIGS: PerDialect<KyselyConfig> = {
 
   mysql: {
     dialect: new MysqlDialect({
+      createConnection,
       pool: async () => createPool(DIALECT_CONFIGS.mysql),
     }),
     plugins: PLUGINS,