feat(live-messages): follow PostgreSQL, Redis and Kafka messages live

apple/Plugins/KafkaDriverPlugin/KafkaPluginDriver+LiveMessages.swift

@@ -0,0 +1,140 @@
+import Foundation
+import os
+import TableProPluginKit
+
+// Follows topics from each partition's end offset with plain Fetch requests and no consumer group, so it never
+// joins a group or moves a group's committed offsets.
+extension KafkaPluginDriver: PluginLiveMessaging {
+    var liveMessageKind: PluginLiveMessageKind? { .topics }
+
+    func subscribe(
+        _ request: PluginLiveSubscriptionRequest,
+        sessions: any PluginLiveSessionSource
+    ) async throws -> any PluginLiveSubscription {
+        guard let follower = try await sessions.openSession() as? KafkaPluginDriver else {
+            throw KafkaError.notConnected
+        }
+        let cluster = follower.liveCluster
+        var cursors: [KafkaFollowCursor] = []
+        for topicName in request.channels {
+            let metadata = try await cluster.metadata(topics: [topicName], refresh: true)
+            let topic = try metadata.requireTopic(named: topicName)
+            let partitions = topic.partitions.map(\.index).sorted()
+            let ends = try await KafkaOffsetsRequest.listOffsets(
+                topic: topicName,
+                partitions: partitions,
+                timestamp: KafkaOffsetsRequest.latestTimestamp,
+                cluster: cluster
+            )
+            cursors += KafkaFollowCursor.cursors(topic: topicName, ends: ends)
+        }
+        return KafkaLiveSubscription(cluster: cluster, cursors: cursors)
+    }
+}
+
+struct KafkaFollowCursor: Sendable, Equatable {
+    let topic: String
+    let partition: Int32
+    let nextOffset: Int64
+
+    static func cursors(topic: String, ends: [Int32: Int64]) -> [KafkaFollowCursor] {
+        ends.keys.sorted().compactMap { partition in
+            ends[partition].map { KafkaFollowCursor(topic: topic, partition: partition, nextOffset: $0) }
+        }
+    }
+
+    static func message(from record: KafkaRecord, topic: String) -> PluginLiveMessage {
+        PluginLiveMessage(
+            timestamp: Date(timeIntervalSince1970: TimeInterval(record.timestamp) / 1_000),
+            channel: topic,
+            payload: text(record.value) ?? "",
+            key: text(record.key),
+            partition: Int(record.partition),
+            offset: record.offset
+        )
+    }
+
+    // A payload that is not UTF-8 reads as hex, the way the topic grid shows it.
+    static func text(_ data: Data?) -> String? {
+        guard let data else { return nil }
+        return String(data: data, encoding: .utf8) ?? data.map { String(format: "%02X", $0) }.joined()
+    }
+}
+
+final class KafkaLiveSubscription: PluginLiveSubscription, @unchecked Sendable {
+    private static let recordsPerFetch = 500
+    private static let idleWait: Duration = .milliseconds(500)
+    private static let logger = Logger(subsystem: "com.TablePro.KafkaDriver", category: "LiveMessages")
+
+    let messages: AsyncThrowingStream<PluginLiveMessage, Error>
+    private let cluster: KafkaCluster
+    private let reader: Task<Void, Never>
+
+    init(cluster: KafkaCluster, cursors: [KafkaFollowCursor]) {
+        self.cluster = cluster
+        let (stream, continuation) = AsyncThrowingStream<PluginLiveMessage, Error>.makeStream()
+        messages = stream
+        reader = Task {
+            do {
+                // Each Fetch waits up to 500 ms on the broker when the partition has nothing new, so no loop spins.
+                try await withThrowingTaskGroup(of: Void.self) { group in
+                    for cursor in cursors {
+                        group.addTask {
+                            var offset = cursor.nextOffset
+                            while !Task.isCancelled {
+                                let started = ContinuousClock.now
+                                let page = try await KafkaFetchRequest.fetch(
+                                    topic: cursor.topic,
+                                    partition: cursor.partition,
+                                    startOffset: offset,
+                                    maximumRecords: Self.recordsPerFetch,
+                                    cluster: cluster
+                                )
+                                for record in page.records {
+                                    continuation.yield(KafkaFollowCursor.message(from: record, topic: cursor.topic))
+                                }
+                                if let last = page.records.last {
+                                    offset = last.offset + 1
+                                } else if ContinuousClock.now - started < Self.idleWait / 2 {
+                                    // A trailing transaction marker answers at once and reads as no record,
+                                    // where an idle partition makes the broker wait out the long poll.
+                                    try await Task.sleep(for: Self.idleWait)
+                                }
+                            }
+                        }
+                    }
+                    // The first partition that fails ends the follow, and leaving the group cancels the others.
+                    while try await group.next() != nil {}
+                }
+                continuation.finish()
+            } catch {
+                guard !Task.isCancelled else {
+                    continuation.finish()
+                    return
+                }
+                Self.logger.info("Topic follow ended: \(String(describing: type(of: error)), privacy: .public) \(error.localizedDescription, privacy: .private)")
+                continuation.finish(throwing: error)
+            }
+        }
+    }
+
+    func publish(_ payload: String, to channel: String) async throws {
+        let metadata = try await cluster.metadata(topics: [channel])
+        let topic = try metadata.requireTopic(named: channel)
+        guard let partition = topic.partitions.map(\.index).min() else { throw KafkaError.unknownTopic(channel) }
+        _ = try await KafkaProduceRequest.produce(
+            topic: channel,
+            partition: partition,
+            key: nil,
+            value: Data(payload.utf8),
+            headers: [],
+            timestamp: Int64(Date().timeIntervalSince1970 * 1_000),
+            cluster: cluster
+        )
+    }
+
+    func close() async {
+        reader.cancel()
+        _ = await reader.value
+    }
+}

apple/Plugins/KafkaDriverPlugin/KafkaPluginDriver.swift

@@ -13,6 +13,8 @@ final class KafkaPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
     private let state = KafkaDriverState()
     private let anchors = KafkaAnchorStore()
 
+    var liveCluster: KafkaCluster { cluster }
+
     var capabilities: PluginCapabilities { [] }
 
     /// The config is read here and not retained. It carries the password, and nothing after

apple/Plugins/PostgreSQLDriverPlugin/LibPQDriverCore.swift

@@ -304,6 +304,11 @@ final class LibPQDriverCore: @unchecked Sendable {
         libpqConnection?.observeServerNotices(observer)
     }
 
+    func nextNotifications(waitMicroseconds: Int64) async throws -> [LibPQNotification] {
+        guard let connection = libpqConnection else { throw LibPQPluginError.notConnected }
+        return try await connection.nextNotifications(waitMicroseconds: waitMicroseconds)
+    }
+
     func cancelQueryKeepingReply() {
         libpqConnection?.requestServerCancelKeepingReply()
     }

apple/Plugins/PostgreSQLDriverPlugin/LibPQPluginConnection.swift

@@ -628,6 +628,39 @@ final class LibPQPluginConnection: @unchecked Sendable {
         typeNames.merge(names)
     }
 
+    // MARK: - Notifications
+
+    // Waits at most `waitMicroseconds` on the queue, so a NOTIFY sent on this session gets its turn between waits.
+    func nextNotifications(waitMicroseconds: Int64) async throws -> [LibPQNotification] {
+        try await pluginDispatchAsync(on: queue) { [self] in
+            guard !isShuttingDown, let conn = connectionHandle else { throw LibPQPluginError.notConnected }
+            var received = Self.drainNotifications(conn)
+            guard received.isEmpty else { return received }
+            let socket = PQsocket(conn)
+            guard socket >= 0 else { throw connectionError(from: conn) }
+            let deadline = pg_usec_time_t(Int64(PQgetCurrentTimeUSec()) + waitMicroseconds)
+            guard PQsocketPoll(socket, 1, 0, deadline) >= 0 else { throw connectionError(from: conn) }
+            guard PQconsumeInput(conn) != 0, PQstatus(conn) == CONNECTION_OK else {
+                stateLock.withLock { _hasLostConnection = true }
+                throw connectionError(from: conn)
+            }
+            received = Self.drainNotifications(conn)
+            return received
+        }
+    }
+
+    private static func drainNotifications(_ conn: OpaquePointer) -> [LibPQNotification] {
+        var received: [LibPQNotification] = []
+        while let notify = PQnotifies(conn) {
+            received.append(LibPQNotification(
+                channel: String(cString: notify.pointee.relname),
+                payload: notify.pointee.extra.map { String(cString: $0) } ?? ""
+            ))
+            PQfreemem(notify)
+        }
+        return received
+    }
+
     // MARK: - Query Cancellation
 
     func cancelCurrentQuery() {

apple/Plugins/PostgreSQLDriverPlugin/PGlitePluginDriver.swift

@@ -12,7 +12,7 @@
 import Foundation
 import TableProPluginKit
 
 final class PGlitePluginDriver: PostgreSQLPluginDriver {
     private let connectHost: String
     private let connectPort: Int
 
@@ -28,6 +28,8 @@
 
     // pldebugger is a server C extension and needs a second session; PGlite has neither.
     override var canDebugRoutines: Bool { false }
+    // A second session would open a different, empty database, so a NOTIFY from the person's session never arrives.
+    override var canFollowNotifications: Bool { false }
 
     override func connect() async throws {
         do {

apple/Plugins/PostgreSQLDriverPlugin/PostgreSQLPluginDriver+LiveMessages.swift

@@ -0,0 +1,96 @@
+//
+//  PostgreSQLPluginDriver+LiveMessages.swift
+//  PostgreSQLDriverPlugin
+//
+
+import Foundation
+import os
+import TableProLogRedaction
+import TableProPluginKit
+
+struct LibPQNotification: Sendable, Equatable {
+    let channel: String
+    let payload: String
+}
+
+// Redshift and CockroachDB are siblings of this class and never conform; PGlite turns it off through
+// `canFollowNotifications`.
+extension PostgreSQLPluginDriver: PluginLiveMessaging {
+    var liveMessageKind: PluginLiveMessageKind? {
+        canFollowNotifications ? .notifications : nil
+    }
+
+    func subscribe(
+        _ request: PluginLiveSubscriptionRequest,
+        sessions: any PluginLiveSessionSource
+    ) async throws -> any PluginLiveSubscription {
+        guard canFollowNotifications,
+              let listener = try await sessions.openSession() as? PostgreSQLPluginDriver
+        else { throw LibPQPluginError.notConnected }
+        listener.core.interruptBlockedCallOnDisconnect()
+        for statement in PostgreSQLListenStatements.listen(request.channels) {
+            _ = try await listener.execute(query: statement)
+        }
+        return PostgreSQLLiveSubscription(listener: listener)
+    }
+}
+
+enum PostgreSQLListenStatements {
+    static let notifyStatement = "SELECT pg_catalog.pg_notify($1, $2)"
+
+    // LISTEN takes an identifier, never a parameter, so each channel is quoted as one.
+    static func listen(_ channels: [String]) -> [String] {
+        channels.map { "LISTEN \(PostgreSQLObjectQueries.quoteIdentifier($0))" }
+    }
+}
+
+final class PostgreSQLLiveSubscription: PluginLiveSubscription, @unchecked Sendable {
+    // Short, so a NOTIFY sent on this session runs between two waits instead of behind a long one.
+    private static let waitMicroseconds: Int64 = 250_000
+    private static let logger = Logger(subsystem: "com.TablePro.PostgreSQLDriver", category: "LiveMessages")
+
+    let messages: AsyncThrowingStream<PluginLiveMessage, Error>
+    private let listener: PostgreSQLPluginDriver
+    private let reader: Task<Void, Never>
+
+    init(listener: PostgreSQLPluginDriver) {
+        self.listener = listener
+        let (stream, continuation) = AsyncThrowingStream<PluginLiveMessage, Error>.makeStream()
+        messages = stream
+        reader = Task {
+            do {
+                while !Task.isCancelled {
+                    let batch = try await listener.core.nextNotifications(waitMicroseconds: Self.waitMicroseconds)
+                    let arrived = Date()
+                    for notification in batch {
+                        continuation.yield(PluginLiveMessage(
+                            timestamp: arrived,
+                            channel: notification.channel,
+                            payload: notification.payload
+                        ))
+                    }
+                }
+                continuation.finish()
+            } catch {
+                guard !Task.isCancelled else {
+                    continuation.finish()
+                    return
+                }
+                Self.logger.info("Notification read ended: \(LogRedaction.publicDescription(of: error), privacy: .public) \(error.localizedDescription, privacy: .private)")
+                continuation.finish(throwing: error)
+            }
+        }
+    }
+
+    func publish(_ payload: String, to channel: String) async throws {
+        _ = try await listener.executeParameterized(
+            query: PostgreSQLListenStatements.notifyStatement,
+            parameters: [.text(channel), .text(payload)]
+        )
+    }
+
+    func close() async {
+        reader.cancel()
+        _ = await reader.value
+    }
+}

apple/Plugins/PostgreSQLDriverPlugin/PostgreSQLPluginDriver.swift

@@ -56,6 +56,7 @@ class PostgreSQLPluginDriver: LibPQBackedDriver, @unchecked Sendable {
 
     // In the class body so PGlite, which inherits the debugger conformance, can turn it off.
     var canDebugRoutines: Bool { true }
+    var canFollowNotifications: Bool { true }
 
     init(config: DriverConnectionConfig, singleConnectionMode: Bool = false) {
         self.core = LibPQDriverCore(config: config, singleConnectionMode: singleConnectionMode)

apple/Plugins/PostgreSQLDriverTests/PostgreSQLListenStatementsTests.swift

@@ -0,0 +1,18 @@
+import Foundation
+import TableProPluginKit
+import Testing
+
+struct PostgreSQLListenStatementsTests {
+    @Test("Each channel is LISTENed as a quoted identifier, so case and quotes survive")
+    func listenQuotesIdentifiers() {
+        #expect(PostgreSQLListenStatements.listen(["Orders", "a\"b"]) == [
+            "LISTEN \"Orders\"",
+            "LISTEN \"a\"\"b\"",
+        ])
+    }
+
+    @Test("NOTIFY goes through pg_notify with bound parameters, never a spliced payload")
+    func notifyIsParameterized() {
+        #expect(PostgreSQLListenStatements.notifyStatement == "SELECT pg_catalog.pg_notify($1, $2)")
+    }
+}

apple/Plugins/RedisDriverPlugin/RedisClusterChannel.swift

@@ -76,6 +76,11 @@ final class RedisClusterChannel: RedisCommandChannel, @unchecked Sendable {
 
     var supportsDatabaseSelection: Bool { servedDatabaseCount > 1 }
 
+    /// Any node: the cluster bus delivers every classic PUBLISH to every node's subscribers.
+    func pubSubConnection() -> RedisPluginConnection? {
+        lock.withLock { connections.values.lazy.compactMap { $0 as? RedisPluginConnection }.first }
+    }
+
     private var servedDatabaseCount: Int {
         lock.lock()
         defer { lock.unlock() }

apple/Plugins/RedisDriverPlugin/RedisCommandChannel.swift

@@ -85,9 +85,14 @@ protocol RedisCommandChannel: AnyObject, Sendable {
     /// Sentinel needs this because a demoted primary keeps answering `role:master` and keeps
     /// accepting writes for several seconds after the quorum has moved on.
     func verifyStillPrimary() async throws
+
+    /// The one node connection a dedicated driver subscribes on; nil when the channel has none.
+    func pubSubConnection() -> RedisPluginConnection?
 }
 
 extension RedisCommandChannel {
+    func pubSubConnection() -> RedisPluginConnection? { nil }
+
     var supportsDatabaseSelection: Bool { true }
     var supportsTransactions: Bool { true }
     var partitionsKeyspace: Bool { false }