...

Author: igalklebanov

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

@@ -42,6 +42,7 @@ export interface MysqlPool {
 }
 
 export interface MysqlPoolConnection {
+  destroy(): void
   query(
     sql: string,
     parameters: ReadonlyArray<unknown>,

src/dialect/mysql/mysql-driver.ts

@@ -60,18 +60,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,
@@ -158,6 +146,18 @@ 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 {
@@ -166,7 +166,7 @@ function isOkPacket(obj: unknown): obj is MysqlOkPacket {
 
 class MysqlConnection implements DatabaseConnection {
   readonly #rawConnection: MysqlPoolConnection
-  #pid: unknown
+  #cid: unknown
 
   constructor(rawConnection: MysqlPoolConnection) {
     this.#rawConnection = rawConnection
@@ -175,15 +175,28 @@ class MysqlConnection implements DatabaseConnection {
   async cancelQuery(
     controlConnectionProvider: ControlConnectionProvider,
   ): Promise<void> {
-    if (!this.#pid) {
+    if (!this.#cid) {
       return logOnce(
         'kysely:warning: cannot cancel query because the connection has not been made cancelable.',
       )
     }
 
-    return await controlConnectionProvider(async (controlConnection) => {
+    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
+    }
+
+    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.#pid]),
+        CompiledQuery.raw('kill connection ?', [this.#cid]),
       )
     })
   }
@@ -284,14 +297,14 @@ class MysqlConnection implements DatabaseConnection {
   }
 
   async #setupCancelability(): Promise<void> {
-    if (this.#pid) {
+    if (this.#cid) {
       return
     }
 
     const results = await this.#executeQuery(
-      CompiledQuery.raw('select connection_id() as pid', []),
+      CompiledQuery.raw('select connection_id() as cid', []),
     )
 
-    this.#pid = (results as any[])[0]?.pid
+    this.#cid = (results as any[])[0]?.cid
   }
 }