78 lines
3.1 KiB
Swift
78 lines
3.1 KiB
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
|
|
}
|
|
}
|