From 614d3ec27cdef50cfb9fc3cfd382b6a4d9578cff Mon Sep 17 00:00:00 2001
From: Tim Condon <0xTim@users.noreply.github.com>
Date: Sun, 11 Aug 2024 09:53:16 +0100
Subject: [PATCH] Add async APIs for shutting down the databases (#614)

---
 Sources/FluentKit/Database/Database.swift  |  7 +++++++
 Sources/FluentKit/Database/Databases.swift | 15 +++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/Sources/FluentKit/Database/Database.swift b/Sources/FluentKit/Database/Database.swift
index 12711d81..f9c2c985 100644
--- a/Sources/FluentKit/Database/Database.swift
+++ b/Sources/FluentKit/Database/Database.swift
@@ -57,6 +57,13 @@ extension Database {
 public protocol DatabaseDriver: Sendable {
     func makeDatabase(with context: DatabaseContext) -> any Database
     func shutdown()
+    func shutdownAsync() async
+}
+
+public extension DatabaseDriver {
+    func shutdownAsync() async {
+        shutdown()
+    }
 }
 
 public protocol DatabaseConfiguration: Sendable {
diff --git a/Sources/FluentKit/Database/Databases.swift b/Sources/FluentKit/Database/Databases.swift
index 3978dcef..3ab006c4 100644
--- a/Sources/FluentKit/Database/Databases.swift
+++ b/Sources/FluentKit/Database/Databases.swift
@@ -141,6 +141,7 @@ public final class Databases: @unchecked Sendable { // @unchecked is safe here;
         self.lock.withLock { Set(self.configurations.keys) }
     }
 
+    @available(*, noasync, message: "Drivers may call wait() and should not be used in an async context", renamed: "shutdownAsync()")
     public func shutdown() {
         self.lock.withLockVoid {
             for driver in self.drivers.values {
@@ -149,6 +150,20 @@ public final class Databases: @unchecked Sendable { // @unchecked is safe here;
             self.drivers = [:]
         }
     }
+    
+    public func shutdownAsync() async {
+        var driversToShutdown: [any DatabaseDriver] = []
+        
+        self.lock.withLockVoid {
+            for driver in self.drivers.values {
+                driversToShutdown.append(driver)
+            }
+            self.drivers = [:]
+        }
+        for driver in driversToShutdown {
+            await driver.shutdownAsync()
+        }
+    }
 
     private func _requireConfiguration(for id: DatabaseID) -> any DatabaseConfiguration {
         guard let configuration = self.configurations[id] else {