feat(connections): pass connection details to scripts, add Redshift IAM sign-in

apple/Packages/TableProCore/Sources/TableProImport/Bundle/ExportableConnection.swift

@@ -143,6 +143,7 @@ public extension ExportableConnection {
     private static let queryTimeoutSecondsRange = 0 ... Int(Int32.max) / 1_000
 
     static let importBlockedAdditionalFieldKeys: Set<String> = [
+        "postdisconnectscript",
         "preconnectscript",
         "pretunnelhost",
         "pretunnelport",

apple/Packages/TableProCore/Tests/TableProImportTests/ShareableAdditionalFieldsTests.swift

@@ -6,6 +6,7 @@ struct ShareableAdditionalFieldsTests {
     @Test func testLeavesOutEveryImportBlockedKeyInAnyCase() {
         let fields = [
             "preConnectScript": "export PGTOKEN=secret-token",
+            "postDisconnectScript": "networksetup -disconnectpppoeservice corp",
             "usePgpass": "true",
             "PROMPTFORPASSWORD": "true",
             "awsRegion": "us-east-1",

apple/TablePro/AppDelegate.swift

@@ -163,6 +163,7 @@ class AppDelegate: NSObject, NSApplicationDelegate {
 
         Task {
             await MCPServerManager.shared.stop()
+            await DatabaseManager.shared.runPostDisconnectScriptsBeforeQuit()
             NSApp.reply(toApplicationShouldTerminate: true)
         }
         return .terminateLater

apple/TablePro/Core/Database/AWS/AWSDeadlineSession.swift

@@ -0,0 +1,26 @@
+//
+//  AWSDeadlineSession.swift
+//  TablePro
+//
+
+import Foundation
+import TableProPluginKit
+
+enum AWSDeadlineSession {
+    // A connect's AWS requests share its deadline; a later sign-in has none and uses the shared session.
+    static func run<Value>(
+        within deadline: ConnectionDeadline?,
+        _ body: (URLSession) async throws -> Value
+    ) async throws -> Value {
+        guard let deadline else {
+            return try await body(AWSHTTP.shared)
+        }
+        let timeout = ConnectionCredentialResolver.credentialRequestTimeout(for: deadline)
+        let configuration = URLSessionConfiguration.default
+        configuration.timeoutIntervalForRequest = timeout
+        configuration.timeoutIntervalForResource = timeout
+        let session = URLSession(configuration: configuration)
+        defer { session.invalidateAndCancel() }
+        return try await body(session)
+    }
+}

apple/TablePro/Core/Database/AWS/AWSIAMTokenSigner.swift

@@ -12,21 +12,10 @@ struct AWSIAMTokenSigner: Sendable {
     let target: Target
 
     func token(deadline: ConnectionDeadline?) async throws -> String {
-        let credentials = try await resolveCredentials(deadline: deadline)
-        return try target.token(username: username, fields: fields, credentials: credentials)
-    }
-
-    private func resolveCredentials(deadline: ConnectionDeadline?) async throws -> AWSCredentials {
-        guard let deadline else {
-            return try await AWSCredentialResolver.resolve(source: source, fields: fields)
+        let credentials = try await AWSDeadlineSession.run(within: deadline) { session in
+            try await AWSCredentialResolver.resolve(source: source, fields: fields, session: session)
         }
-        let timeout = ConnectionCredentialResolver.credentialRequestTimeout(for: deadline)
-        let configuration = URLSessionConfiguration.default
-        configuration.timeoutIntervalForRequest = timeout
-        configuration.timeoutIntervalForResource = timeout
-        let session = URLSession(configuration: configuration)
-        defer { session.invalidateAndCancel() }
-        return try await AWSCredentialResolver.resolve(source: source, fields: fields, session: session)
+        return try target.token(username: username, fields: fields, credentials: credentials)
     }
 }
 

apple/TablePro/Core/Database/AWS/AWSJSONRequest.swift

@@ -0,0 +1,100 @@
+//
+//  AWSJSONRequest.swift
+//  TablePro
+//
+
+import Foundation
+import TableProPluginKit
+
+// A SigV4-signed call to an AWS service that speaks the JSON 1.1 protocol.
+struct AWSJSONRequest: Sendable, Equatable {
+    static let contentType = "application/x-amz-json-1.1"
+
+    let service: String
+    let region: String
+    let host: String
+    let target: String
+    let body: Data
+
+    init(service: String, region: String, target: String, body: Data) {
+        let canonicalRegion = AWSPartition.canonicalRegion(region)
+        self.service = service
+        self.region = canonicalRegion
+        host = AWSPartition.host(service: service, region: canonicalRegion)
+        self.target = target
+        self.body = body
+    }
+
+    func signedURLRequest(credentials: AWSCredentials, now: Date = Date()) throws -> URLRequest {
+        guard let url = URL(string: "https://\(host)/") else {
+            throw AWSAuthError.missingConfiguration(
+                String(
+                    format: String(localized: "Could not build the %1$@ endpoint for the AWS region \"%2$@\"."),
+                    service, region
+                )
+            )
+        }
+
+        let stamps = Self.timestamps(for: now)
+        var headers = [
+            "content-type": Self.contentType,
+            "host": host,
+            "x-amz-date": stamps.amzDate,
+            "x-amz-target": target
+        ]
+        if let sessionToken = credentials.sessionToken, !sessionToken.isEmpty {
+            headers["x-amz-security-token"] = sessionToken
+        }
+
+        let signedHeaderNames = headers.keys.sorted()
+        let canonicalHeaders = signedHeaderNames.map { "\($0):\(headers[$0] ?? "")\n" }.joined()
+        let signedHeaders = signedHeaderNames.joined(separator: ";")
+        let canonicalRequest = [
+            "POST",
+            "/",
+            "",
+            canonicalHeaders,
+            signedHeaders,
+            AWSSigV4.sha256Hex(body)
+        ].joined(separator: "\n")
+
+        let credentialScope = "\(stamps.dateStamp)/\(region)/\(service)/aws4_request"
+        let stringToSign = [
+            "AWS4-HMAC-SHA256",
+            stamps.amzDate,
+            credentialScope,
+            AWSSigV4.sha256Hex(Data(canonicalRequest.utf8))
+        ].joined(separator: "\n")
+
+        let signingKey = AWSSigV4.deriveSigningKey(
+            secretKey: credentials.secretAccessKey,
+            dateStamp: stamps.dateStamp,
+            region: region,
+            service: service
+        )
+        let signature = AWSSigV4.hmacHex(key: signingKey, data: Data(stringToSign.utf8))
+
+        var request = URLRequest(url: url)
+        request.httpMethod = "POST"
+        request.httpBody = body
+        for name in signedHeaderNames {
+            request.setValue(headers[name], forHTTPHeaderField: name)
+        }
+        request.setValue(
+            "AWS4-HMAC-SHA256 Credential=\(credentials.accessKeyId)/\(credentialScope), "
+                + "SignedHeaders=\(signedHeaders), Signature=\(signature)",
+            forHTTPHeaderField: "Authorization"
+        )
+        return request
+    }
+
+    private static func timestamps(for now: Date) -> (amzDate: String, dateStamp: String) {
+        let formatter = DateFormatter()
+        formatter.locale = Locale(identifier: "en_US_POSIX")
+        formatter.timeZone = TimeZone(identifier: "UTC")
+        formatter.dateFormat = "yyyyMMdd'T'HHmmss'Z'"
+        let amzDate = formatter.string(from: now)
+        formatter.dateFormat = "yyyyMMdd"
+        return (amzDate, formatter.string(from: now))
+    }
+}