From ae5cb1e297680718ceb5f3727feedebe9fb2a50b Mon Sep 17 00:00:00 2001
From: Igal Klebanov <igalklebanov@gmail.com>
Date: Sat, 19 Oct 2024 16:32:22 +0300
Subject: [PATCH] `await using kysely = new Kysely()` support. (#1167)

---
 src/kysely.ts                       |  9 +++-
 test/node/src/async-dispose.test.ts | 65 +++++++++++++++++++++++++++++
 test/node/tsconfig.json             |  2 +
 3 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 test/node/src/async-dispose.test.ts

diff --git a/src/kysely.ts b/src/kysely.ts
index 030e2d4be..b5f84b60d 100644
--- a/src/kysely.ts
+++ b/src/kysely.ts
@@ -42,6 +42,9 @@ import {
   provideControlledConnection,
 } from './util/provide-controlled-connection.js'
 
+// @ts-ignore
+Symbol.asyncDispose ??= Symbol('Symbol.asyncDispose')
+
 /**
  * The main Kysely class.
  *
@@ -87,7 +90,7 @@ import {
  */
 export class Kysely<DB>
   extends QueryCreator<DB>
-  implements QueryExecutorProvider
+  implements QueryExecutorProvider, AsyncDisposable
 {
   readonly #props: KyselyProps
 
@@ -509,6 +512,10 @@ export class Kysely<DB>
 
     return this.getExecutor().executeQuery<R>(compiledQuery, queryId)
   }
+
+  async [Symbol.asyncDispose]() {
+    await this.destroy()
+  }
 }
 
 export class Transaction<DB> extends Kysely<DB> {
diff --git a/test/node/src/async-dispose.test.ts b/test/node/src/async-dispose.test.ts
new file mode 100644
index 000000000..72440ba46
--- /dev/null
+++ b/test/node/src/async-dispose.test.ts
@@ -0,0 +1,65 @@
+import {
+  CompiledQuery,
+  DatabaseConnection,
+  DummyDriver,
+  Kysely,
+  PostgresAdapter,
+  PostgresIntrospector,
+  PostgresQueryCompiler,
+  QueryResult,
+  RootOperationNode,
+  sql,
+} from '../../..'
+import { expect } from './test-setup'
+
+describe('async dispose', function () {
+  it('should call destroy ', async () => {
+    const steps: string[] = []
+
+    {
+      await using db = new Kysely({
+        dialect: {
+          createAdapter: () => new PostgresAdapter(),
+          createDriver: () =>
+            new (class extends DummyDriver {
+              async acquireConnection() {
+                return new (class implements DatabaseConnection {
+                  async executeQuery<R>(): Promise<QueryResult<R>> {
+                    steps.push('executed')
+                    return { rows: [] }
+                  }
+                  streamQuery<R>(): AsyncIterableIterator<QueryResult<R>> {
+                    throw new Error('Method not implemented.')
+                  }
+                })()
+              }
+              async destroy(): Promise<void> {
+                steps.push('destroyed')
+              }
+            })(),
+          createIntrospector: (db) => new PostgresIntrospector(db),
+          createQueryCompiler: () =>
+            new (class extends PostgresQueryCompiler {
+              compileQuery(node: RootOperationNode): CompiledQuery<unknown> {
+                const compiled = super.compileQuery(node)
+                steps.push('compiled')
+                return compiled
+              }
+            })(),
+        },
+      })
+
+      await sql`select 1`.execute(db)
+    }
+
+    steps.push('after runScope')
+
+    expect(steps).to.length.to.be.greaterThan(1)
+    expect(steps).to.deep.equal([
+      'compiled',
+      'executed',
+      'destroyed',
+      'after runScope',
+    ])
+  })
+})
diff --git a/test/node/tsconfig.json b/test/node/tsconfig.json
index 83e6196bb..1ed9e5376 100644
--- a/test/node/tsconfig.json
+++ b/test/node/tsconfig.json
@@ -2,6 +2,8 @@
   "extends": "../../tsconfig-base.json",
   "include": ["src/**/*"],
   "compilerOptions": {
+    "target": "ES2022",
+    "lib": ["ESNext"],
     "module": "CommonJS",
     "outDir": "dist",
     "skipLibCheck": true