From fb5bf4d75846261da93a5480cdefafefdd3df19f Mon Sep 17 00:00:00 2001
From: Igal Klebanov <igalklebanov@gmail.com>
Date: Tue, 18 Jun 2024 12:16:11 +0300
Subject: [PATCH] handle connection errors @ `MssqlDriver`. (#1042)

---
 src/dialect/mssql/mssql-dialect-config.ts | 20 +++++----
 src/dialect/mssql/mssql-driver.ts         | 32 ++++++++++----
 test/node/src/disconnects.test.ts         | 52 +++++++++++++++++++++++
 3 files changed, 88 insertions(+), 16 deletions(-)
 create mode 100644 test/node/src/disconnects.test.ts

diff --git a/src/dialect/mssql/mssql-dialect-config.ts b/src/dialect/mssql/mssql-dialect-config.ts
index e84528f04..6adc2f09d 100644
--- a/src/dialect/mssql/mssql-dialect-config.ts
+++ b/src/dialect/mssql/mssql-dialect-config.ts
@@ -62,22 +62,26 @@ export interface TediousConnection {
     name?: string,
     isolationLevel?: number,
   ): void
+  cancel(): boolean
+  close(): void
   commitTransaction(
     callback: (error?: Error | null) => void,
     name?: string,
   ): void
+  connect(callback?: (error?: Error) => void): void
   execSql(request: TediousRequest): void
+  off(event: 'error', listener: (error: unknown) => void): this
+  off(event: string, listener: (...args: any[]) => void): this
+  on(event: 'error', listener: (error: unknown) => void): this
+  on(event: string, listener: (...args: any[]) => void): this
+  once(event: 'end', listener: () => void): this
+  once(event: string, listener: (...args: any[]) => void): this
+  reset(callback: (error?: Error | null) => void): void
   rollbackTransaction(
     callback: (error?: Error | null) => void,
     name?: string,
   ): void
   saveTransaction(callback: (error?: Error | null) => void, name: string): void
-  cancel(): boolean
-  reset(callback: (error?: Error | null) => void): void
-  close(): void
-  once(event: 'end', listener: () => void): this
-  once(event: string, listener: (...args: any[]) => void): this
-  connect(callback?: (error?: Error) => void): void
 }
 
 export type TediousIsolationLevel = Record<string, number>
@@ -104,12 +108,12 @@ export declare class TediousRequest {
       scale?: number
     }> | null,
   ): void
+  off(event: 'row', listener: (columns: any) => void): this
+  off(event: string, listener: (...args: any[]) => void): this
   on(event: 'row', listener: (columns: any) => void): this
   on(event: string, listener: (...args: any[]) => void): this
   once(event: 'requestCompleted', listener: () => void): this
   once(event: string, listener: (...args: any[]) => void): this
-  off(event: 'row', listener: (columns: any) => void): this
-  off(event: string, listener: (...args: any[]) => void): this
   pause(): void
   resume(): void
 }
diff --git a/src/dialect/mssql/mssql-driver.ts b/src/dialect/mssql/mssql-driver.ts
index ec749f685..71807b37c 100644
--- a/src/dialect/mssql/mssql-driver.ts
+++ b/src/dialect/mssql/mssql-driver.ts
@@ -46,14 +46,10 @@ export class MssqlDriver implements Driver {
       create: async () => {
         const connection = await this.#config.tedious.connectionFactory()
 
-        await new Promise((resolve, reject) =>
-          connection.connect((error) => {
-            if (error) reject(error)
-            else resolve(undefined)
-          }),
-        )
-
-        return new MssqlConnection(connection, this.#config.tedious)
+        return await new MssqlConnection(
+          connection,
+          this.#config.tedious,
+        ).connect()
       },
       destroy: async (connection) => {
         await connection[PRIVATE_DESTROY_METHOD]()
@@ -104,6 +100,11 @@ class MssqlConnection implements DatabaseConnection {
   constructor(connection: TediousConnection, tedious: Tedious) {
     this.#connection = connection
     this.#tedious = tedious
+
+    this.#connection.on('error', console.error)
+    this.#connection.once('end', () => {
+      this.#connection.off('error', console.error)
+    })
   }
 
   async beginTransaction(settings: TransactionSettings): Promise<void> {
@@ -132,6 +133,21 @@ class MssqlConnection implements DatabaseConnection {
     )
   }
 
+  async connect(): Promise<this> {
+    await new Promise((resolve, reject) => {
+      this.#connection.connect((error) => {
+        if (error) {
+          console.error(error)
+          reject(error)
+        } else {
+          resolve(undefined)
+        }
+      })
+    })
+
+    return this
+  }
+
   async executeQuery<O>(compiledQuery: CompiledQuery): Promise<QueryResult<O>> {
     try {
       const deferred = new Deferred<OnDone<O>>()
diff --git a/test/node/src/disconnects.test.ts b/test/node/src/disconnects.test.ts
new file mode 100644
index 000000000..56ae1ff36
--- /dev/null
+++ b/test/node/src/disconnects.test.ts
@@ -0,0 +1,52 @@
+import { Kysely, MssqlDialect, sql } from '../../..'
+import { DIALECTS, DIALECT_CONFIGS, Database, expect } from './test-setup'
+import * as tarn from 'tarn'
+import * as tedious from 'tedious'
+
+const dialect = 'mssql'
+
+if (DIALECTS.includes(dialect)) {
+  describe(`${dialect}: disconnects`, () => {
+    let connection: tedious.Connection
+    let connectionFactoryTimesCalled = 0
+    let db: Kysely<Database>
+
+    before(async () => {
+      db = new Kysely({
+        dialect: new MssqlDialect({
+          tarn: {
+            ...tarn,
+            options: {
+              min: 0,
+              max: 1,
+            },
+          },
+          tedious: {
+            ...tedious,
+            connectionFactory: () => {
+              connectionFactoryTimesCalled++
+
+              return (connection = new tedious.Connection(
+                DIALECT_CONFIGS[dialect],
+              ))
+            },
+          },
+        }),
+      })
+    })
+
+    after(async () => {
+      await db.destroy()
+    })
+
+    it('should be disconnection tolerant', async () => {
+      await sql`select 1`.execute(db)
+      expect(connectionFactoryTimesCalled).to.equal(1)
+
+      connection.socketError(new Error('moshe'))
+
+      await sql`select 1`.execute(db)
+      expect(connectionFactoryTimesCalled).to.equal(2)
+    })
+  })
+}