Shutdown the pool on hard errors, add some tests for this

trasch · trasch · commit e326808fbff5 · 2023-05-08T18:21:53.000+02:00
diff --git a/Package.resolved b/Package.resolved
diff --git a/Package.swift b/Package.swift
@@ -12,7 +12,7 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-collections", from: "1.0.2"),
-        .package(url: "https://github.com/vapor/postgres-nio.git", from: "1.14.0"),
+        .package(url: "https://github.com/vapor/postgres-nio.git", from: "1.14.2"),
         .package(url: "https://github.com/vapor/postgres-kit", from: "2.10.0"),
     ],
     targets: [
@@ -23,4 +23,9 @@ let package = Package(
                 .product(name: "PostgresNIO", package: "postgres-nio"),
                 .product(name: "PostgresKit", package: "postgres-kit"),
             ]),
+        .testTarget(
+            name: "PostgresConnectionPoolTests",
+            dependencies: [
+                .target(name: "PostgresConnectionPool"),
+            ]),
     ])
diff --git a/Sources/PostgresConnectionPool/Extensions/PSQLError+Description.swift b/Sources/PostgresConnectionPool/Extensions/PSQLError+Description.swift
@@ -0,0 +1,20 @@
+//
+//  Created by Thomas Rasch on 05.05.23.
+//
+
+import PostgresNIO
+
+extension PSQLError: CustomStringConvertible {
+
+    public var description: String {
+        if let serverInfo = self.serverInfo,
+           let severity = serverInfo[.severity],
+           let message = serverInfo[.message]
+        {
+            return "\(severity): \(message)"
+        }
+
+        return "Database error: \(self.code.description)"
+    }
+
+}
diff --git a/Sources/PostgresConnectionPool/PoolError.swift b/Sources/PostgresConnectionPool/PoolError.swift
@@ -22,3 +22,18 @@ public enum PoolError: Error {
     case unknown
 
 }
+
+extension PoolError: CustomStringConvertible {
+
+    public var description: String {
+        switch self {
+        case .cancelled: return "<PoolError: cancelled>"
+        case .connectionFailed: return "<PoolError: connectionFailed>"
+        case .poolDestroyed: return "<PoolError: poolDestroyed>"
+        case .postgresError(let psqlError): return "<PoolError: postgresError='\(psqlError.description)'>"
+        case .queryCancelled(query: let query, runtime: let runtime): return "<PoolError: query '\(query)' cancelled after \(runtime.rounded(toPlaces: 3))s>"
+        case .unknown: return "<PoolError: unknown>"
+        }
+    }
+
+}
diff --git a/Sources/PostgresConnectionPool/PostgresConnectionPool.swift b/Sources/PostgresConnectionPool/PostgresConnectionPool.swift
@@ -33,7 +33,8 @@ public actor PostgresConnectionPool {
     private var inUseConnectionCounts: Deque<Int> = []
 
     private var didStartWatcherTask = false
-    private var didShutdown = false
+    public private(set) var didShutdown = false
+    public private(set) var shutdownError: PoolError?
 
     // MARK: -
 
@@ -158,6 +159,8 @@ public actor PostgresConnectionPool {
     ///
     /// Must be done here since Swift doesn't yet allow async deinit.
     public func shutdown() async {
+        guard !didShutdown else { return }
+
         logger.debug("[\(poolName)] shutdown()")
 
         didShutdown = true
@@ -181,6 +184,15 @@ public actor PostgresConnectionPool {
         try? await eventLoopGroup.shutdownGracefully()
     }
 
+    // Shutdown due to an unrecoverable error
+    private func shutdown(withError error: PoolError) async {
+        guard !didShutdown else { return }
+
+        shutdownError = error
+
+        await shutdown()
+    }
+
     /// Information about the pool and its open connections.
     func poolInfo() async -> PoolInfo {
         let connections = connections.compactMap { connection -> PoolInfo.ConnectionInfo? in
@@ -438,23 +450,54 @@ public actor PostgresConnectionPool {
             let logMessage: Logger.Message = "[\(poolName)] Connection \(poolConnection.id) failed after \(connectionRuntime.rounded(toPlaces: 2))s: \(error)"
 
             // Don't just open a new connection, check first if this is a permanent error like wrong password etc.
-            if let pslError = error as? PSQLError,
-               let serverInfo = pslError.serverInfo,
-               let code = serverInfo[.sqlState]
-            {
-                // TODO: List of hard errors
-                switch PostgresError.Code(raw: code) {
-                case .adminShutdown,
-                     .cannotConnectNow,
-                     .invalidAuthorizationSpecification,
-                     .invalidName,
-                     .invalidPassword:
+
+            // TODO: Check which of these errors is actually fatal
+            if let psqlError = error as? PSQLError {
+                switch psqlError.code {
+                case .authMechanismRequiresPassword,
+                     .connectionClosed,
+                     .connectionError,
+                     .connectionQuiescing,
+                     .sslUnsupported,
+                     .failedToAddSSLHandler,
+                     .invalidCommandTag,
+                     .messageDecodingFailure,
+                     .receivedUnencryptedDataAfterSSLRequest,
+                     .saslError,
+                     .server,
+                     .unexpectedBackendMessage,
+                     .unsupportedAuthMechanism:
                     logger.error(logMessage)
+                    await shutdown(withError: PoolError.postgresError(psqlError))
                     return
 
+                case .queryCancelled,
+                     .tooManyParameters,
+                     .uncleanShutdown:
+                    break
+
                 default:
                     break
                 }
+
+                if let serverInfo = psqlError.serverInfo,
+                   let code = serverInfo[.sqlState]
+                {
+                    // TODO: Check which of these errors is actually fatal
+                    switch PostgresError.Code(raw: code) {
+                    case .adminShutdown,
+                         .cannotConnectNow,
+                         .invalidAuthorizationSpecification,
+                         .invalidName,
+                         .invalidPassword:
+                        logger.error(logMessage)
+                        await shutdown(withError: PoolError.postgresError(psqlError))
+                        return
+
+                    default:
+                        break
+                    }
+                }
             }
 
             logger.warning(logMessage)
diff --git a/Tests/PostgresConnectionPoolTests/ConnectionErrorTests.swift b/Tests/PostgresConnectionPoolTests/ConnectionErrorTests.swift
@@ -0,0 +1,97 @@
+//
+//  Created by Thomas Rasch on 05.05.23.
+//
+
+import PostgresConnectionPool
+import PostgresNIO
+import XCTest
+
+final class ConnectionErrorTests: XCTestCase {
+
+    private var logger: Logger = {
+        var logger = Logger(label: "ConnectionErrorTests")
+        logger.logLevel = .info
+        return logger
+    }()
+
+    // TODO: Clean up the error checking
+
+    private func withConfiguration(
+        _ configuration: PoolConfiguration,
+        expectedErrorDescription: String)
+        async throws
+    {
+        let pool = PostgresConnectionPool(configuration: configuration, logger: logger)
+        do {
+            try await pool.connection { connection in
+                try await connection.query("SELECT 1", logger: logger)
+            }
+            await pool.shutdown()
+
+            XCTFail("Can't connect, so we should have an exception")
+        }
+        catch {
+            XCTAssertTrue(error is PoolError)
+
+            let shutdownError = await pool.shutdownError
+            XCTAssertEqual(shutdownError?.description, expectedErrorDescription)
+        }
+        let didShutdown = await pool.didShutdown
+        XCTAssertTrue(didShutdown)
+    }
+
+    func testConnectWrongHost() async throws {
+        try await withConfiguration(self.poolConfiguration(host: "notworking"), expectedErrorDescription: "<PoolError: postgresError='Database error: connectionError'>")
+    }
+
+    func testConnectWrongPort() async throws {
+        try await withConfiguration(self.poolConfiguration(port: 99999), expectedErrorDescription: "<PoolError: postgresError='Database error: connectionError'>")
+    }
+
+    func testConnectWrongUsername() async throws {
+        try await withConfiguration(self.poolConfiguration(username: "notworking"), expectedErrorDescription: "<PoolError: postgresError='FATAL: password authentication failed for user \"notworking\"'>")
+    }
+
+    func testConnectWrongPassword() async throws {
+        try await withConfiguration(self.poolConfiguration(password: "notworking"), expectedErrorDescription: "<PoolError: postgresError='FATAL: password authentication failed for user \"test_username\"'>")
+    }
+
+    func testConnectInvalidTLSConfig() async throws {
+        var tlsConfiguration: TLSConfiguration = .clientDefault
+        tlsConfiguration.maximumTLSVersion = .tlsv1 // New Postgres versions want at least TLSv1.2
+
+        let tls: PostgresConnection.Configuration.TLS = .require(try .init(configuration: tlsConfiguration))
+        try await withConfiguration(self.poolConfiguration(tls: tls), expectedErrorDescription: "<PoolError: postgresError='Database error: sslUnsupported'>")
+    }
+
+    // MARK: -
+
+    private func poolConfiguration(
+        host: String? = nil,
+        port: Int? = nil,
+        username: String? = nil,
+        password: String? = nil,
+        database: String? = nil,
+        tls: PostgresConnection.Configuration.TLS = .disable)
+        -> PoolConfiguration
+    {
+        let postgresConfiguration = PostgresConnection.Configuration(
+            host: host ?? env("POSTGRES_HOSTNAME") ?? "localhost",
+            port: port ?? env("POSTGRES_PORT").flatMap(Int.init(_:)) ?? 5432,
+            username: username ?? env("POSTGRES_USER") ?? "test_username",
+            password: password ?? env("POSTGRES_PASSWORD") ?? "test_password",
+            database: database ?? env("POSTGRES_DB") ?? "test_database",
+            tls: tls)
+        return PoolConfiguration(
+            applicationName: "ConnectionErrorTests",
+            postgresConfiguration: postgresConfiguration,
+            connectTimeout: 10.0,
+            queryTimeout: 10.0,
+            poolSize: 5,
+            maxIdleConnections: 1)
+    }
+
+    private func env(_ name: String) -> String? {
+        getenv(name).flatMap { String(cString: $0) }
+    }
+}
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -0,0 +1,15 @@
+version: '3.7'
+
+x-shared-config: &shared_config
+  environment:
+    POSTGRES_HOST_AUTH_METHOD: "${POSTGRES_HOST_AUTH_METHOD:-scram-sha-256}"
+    POSTGRES_USER: test_username
+    POSTGRES_DB: test_database
+    POSTGRES_PASSWORD: test_password
+  ports:
+    - 5432:5432
+
+services:
+  psql-14:
+    image: postgres:14
+    <<: *shared_config
