...

...

...

...

...

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,10 @@ 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,
     parameters: ReadonlyArray<unknown>,
@@ -51,6 +61,12 @@ export interface MysqlPoolConnection {
     parameters: ReadonlyArray<unknown>,
     callback: (error: unknown, result: MysqlQueryResult) => void,
   ): void
+  threadId: number
+}
+
+export type MysqlConectionConstructor = new (opts?: object) => MysqlConnection
+
+export interface MysqlPoolConnection extends MysqlConnection {
   release(): void
 }
 

src/dialect/mysql/mysql-driver.ts

@@ -1,4 +1,5 @@
 import {
+  ControlConnectionProvider,
   DatabaseConnection,
   QueryResult,
 } from '../../driver/database-connection.js'
@@ -39,7 +40,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
@@ -57,18 +61,6 @@ export class MysqlDriver implements Driver {
     return connection
   }
 
-  async #acquireConnection(): Promise<MysqlPoolConnection> {
-    return new Promise((resolve, reject) => {
-      this.#pool!.getConnection(async (err, rawConnection) => {
-        if (err) {
-          reject(err)
-        } else {
-          resolve(rawConnection)
-        }
-      })
-    })
-  }
-
   async beginTransaction(
     connection: DatabaseConnection,
     settings: TransactionSettings,
@@ -155,106 +147,163 @@ export class MysqlDriver implements Driver {
       })
     })
   }
+
+  async #acquireConnection(): Promise<MysqlPoolConnection> {
+    return new Promise((resolve, reject) => {
+      this.#pool!.getConnection(async (err, rawConnection) => {
+        if (err) {
+          reject(err)
+        } else {
+          resolve(rawConnection)
+        }
+      })
+    })
+  }
 }
 
 function isOkPacket(obj: unknown): obj is MysqlOkPacket {
   return isObject(obj) && 'insertId' in obj && 'affectedRows' in obj
 }
 
 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
   }
 
+  async cancelQuery(
+    controlConnectionProvider: ControlConnectionProvider,
+  ): Promise<void> {
+    try {
+      // this removes the connection from the pool.
+      // this is done to avoid picking it up after the `kill` command next, which
+      // would cause an error when attempting to query.
+      this.#rawConnection.destroy()
+    } catch {
+      // 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) => {
+      await controlConnection.executeQuery(CompiledQuery.raw(cancelQuery, []))
+    })
+  }
+
   async executeQuery<O>(compiledQuery: CompiledQuery): Promise<QueryResult<O>> {
     try {
       const result = await this.#executeQuery(compiledQuery)
 
-      if (isOkPacket(result)) {
-        const { insertId, affectedRows, changedRows } = result
-
-        return {
-          insertId:
-            insertId !== undefined &&
-            insertId !== null &&
-            insertId.toString() !== '0'
-              ? BigInt(insertId)
-              : undefined,
-          numAffectedRows:
-            affectedRows !== undefined && affectedRows !== null
-              ? BigInt(affectedRows)
-              : undefined,
-          numChangedRows:
-            changedRows !== undefined && changedRows !== null
-              ? BigInt(changedRows)
-              : undefined,
-          rows: [],
-        }
-      } else if (Array.isArray(result)) {
+      if (!isOkPacket(result)) {
         return {
-          rows: result as O[],
+          rows: (Array.isArray(result) ? result : []) as never,
         }
       }
 
+      const { insertId, affectedRows, changedRows } = result
+
       return {
+        insertId:
+          insertId !== undefined &&
+          insertId !== null &&
+          insertId.toString() !== '0'
+            ? BigInt(insertId)
+            : undefined,
+        numAffectedRows:
+          affectedRows !== undefined && affectedRows !== null
+            ? BigInt(affectedRows)
+            : undefined,
+        numChangedRows:
+          changedRows !== undefined && changedRows !== null
+            ? BigInt(changedRows)
+            : undefined,
         rows: [],
       }
     } catch (err) {
       throw extendStackTrace(err, new Error())
     }
   }
 
-  #executeQuery(compiledQuery: CompiledQuery): Promise<MysqlQueryResult> {
-    return new Promise((resolve, reject) => {
-      this.#rawConnection.query(
-        compiledQuery.sql,
-        compiledQuery.parameters,
-        (err, result) => {
-          if (err) {
-            reject(err)
-          } else {
-            resolve(result)
-          }
-        },
-      )
-    })
-  }
-
   async *streamQuery<O>(
     compiledQuery: CompiledQuery,
     _chunkSize: number,
   ): AsyncIterableIterator<QueryResult<O>> {
     const stream = this.#rawConnection
       .query(compiledQuery.sql, compiledQuery.parameters)
-      .stream<O>({
-        objectMode: true,
-      })
+      .stream<O>({ objectMode: true })
 
     try {
       for await (const row of stream) {
         yield {
           rows: [row],
         }
       }
-    } catch (ex) {
+    } catch (error) {
       if (
-        ex &&
-        typeof ex === 'object' &&
-        'code' in ex &&
-        // @ts-ignore
-        ex.code === 'ERR_STREAM_PREMATURE_CLOSE'
+        error &&
+        typeof error === 'object' &&
+        'code' in error &&
+        error.code === 'ERR_STREAM_PREMATURE_CLOSE'
       ) {
         // Most likely because of https://github.com/mysqljs/mysql/blob/master/lib/protocol/sequences/Query.js#L220
         return
       }
 
-      throw ex
+      throw error
     }
   }
 
   [PRIVATE_RELEASE_METHOD](): void {
     this.#rawConnection.release()
   }
+
+  #executeQuery(compiledQuery: CompiledQuery): Promise<MysqlQueryResult> {
+    return new Promise((resolve, reject) => {
+      this.#rawConnection.query(
+        compiledQuery.sql,
+        compiledQuery.parameters,
+        (err, result) => {
+          if (err) {
+            reject(err)
+          } else {
+            resolve(result)
+          }
+        },
+      )
+    })
+  }
 }

test/node/src/cancellation.test.ts

@@ -12,7 +12,7 @@ import {
 } from './test-setup.js'
 
 for (const dialect of DIALECTS) {
-  if (dialect === 'postgres') {
+  if (dialect === 'postgres' || dialect === 'mysql') {
     describe(`${dialect}: query cancellation`, () => {
       let ctx: TestContext
 
@@ -64,7 +64,7 @@ for (const dialect of DIALECTS) {
         const delayedQuery = (
           {
             postgres: sql`select pg_sleep(0.1); ${writeQuery};`,
-            mysql: sql`select sleep(0.1); ${writeQuery};`,
+            mysql: sql`do sleep(0.1); ${writeQuery};`,
             mssql: sql`waitfor delay '00:00:00.100'; ${writeQuery};`,
             sqlite: sql`WITH RECURSIVE timer(i) AS (
               SELECT 1

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'
@@ -201,6 +201,7 @@ export const DB_CONFIGS: PerDialectVariant<KyselyConfig> = {
 
   mysql: {
     dialect: new MysqlDialect({
+      createConnection,
       pool: async () => createPool(DIALECT_CONFIGS.mysql),
     }),
     plugins: PLUGINS,