ccn/Packages/Persistence/Sources/Public/Methods/Probe.swift

import HummingbirdFluent
import SQLKit

/// A readiness probe reporting whether the database behind a `Fluent` service is reachable.
///
/// Built once around the application's `Fluent` service and called as a function whenever a fresh answer is needed — typically from a readiness endpoint:
/// `let ready = await probe()`.
public struct Probe: Sendable {

    // MARK: Properties

    /// The `Fluent` service whose default database is probed.
    private let fluent: Fluent

    /// The longest the probe waits for the database's answer before reporting it as not reachable.
    private let timeout: Duration

    // MARK: Initializers

    /// Creates a probe for the default database of the given `Fluent` service.
    /// - Parameters:
    ///   - fluent: the `Fluent` service whose default database is probed.
    ///   - timeout: the longest the probe waits for the database's answer before reporting it as not reachable.
    public init(
        fluent: Fluent,
        timeout: Duration = .seconds(2)
    ) {
        self.fluent = fluent
        self.timeout = timeout
    }

    // MARK: Methods

    /// Reports whether the database behind the `Fluent` service is reachable.
    ///
    /// Runs a trivial `SELECT 1` against the default database — the cheapest statement both the PostgreSQL and SQLite backends understand — so
    /// a readiness check does not depend on any particular schema or model. Any failure (connection refused, authentication error, pool exhausted) is
    /// reported as not reachable rather than thrown, and an answer that does not arrive within the timeout is likewise reported as not reachable — so
    /// a database that hangs yields a prompt "not ready" instead of a hanging readiness endpoint. A default database that is not an SQL database is
    /// also reported as not reachable.
    /// - Returns: `true` when the database answers the probe in time, `false` otherwise.
    public func callAsFunction() async -> Bool {
        guard let database = fluent.db() as? any SQLDatabase else {
            return false
        }

        // The query is raced against the deadline from unstructured tasks whose first answer wins: a structured group
        // would await the query child before returning, so a database that hangs mid-query — the very failure the
        // deadline exists for — would hang the probe with it. The loser is cancelled and abandoned; a late answer lands
        // in a finished stream and is dropped.
        let (answers, continuation) = AsyncStream.makeStream(of: Bool.self)
        let query = Task {
            do {
                try await database
                    .raw("SELECT 1")
                    .run()

                continuation.yield(true)
            } catch {
                continuation.yield(false)
            }
        }
        let deadline = Task {
            try? await Task.sleep(for: timeout)

            continuation.yield(false)
        }

        var answer = answers.makeAsyncIterator()
        let isReachable = await answer.next() ?? false

        query.cancel()
        deadline.cancel()

        return isReachable
    }
}