feat(cassandra): set the consistency level per connection and with CONSISTENCY

apple/Plugins/CassandraDriverPlugin/CassandraConnection.swift

@@ -19,7 +19,7 @@
         return f
     }()
 
     nonisolated(unsafe) private static let dateFormatter: DateFormatter = {
         let f = DateFormatter()
         f.locale = Locale(identifier: "en_US_POSIX")
         f.dateFormat = "yyyy-MM-dd"
@@ -236,25 +236,30 @@
     /// Reads every row a statement returns, a page at a time, so a large result neither stops at a count of the
     /// driver's choosing nor asks the server for more than it answers unpaged: ScyllaDB aborts an unpaged read past
     /// 100 MB. A statement the server will not page, `IN` with `ORDER BY` on the partition key, is read unpaged.
-    func executeQuery(_ cql: String, cancellation: CassandraCancellation? = nil) throws -> CassandraRawResult {
+    func executeQuery(
+        _ cql: String,
+        consistency: CassandraConsistency?,
+        cancellation: CassandraCancellation? = nil
+    ) throws -> CassandraRawResult {
         guard let session else {
             throw CassandraPluginError.notConnected
         }
         forgetResumePointsUnlessRead(cql)
 
         let startTime = Date()
         do {
-            return try executePaged(cql, pageSize: Self.queryPageSize, session: session, cancellation: cancellation,
-                                    startTime: startTime)
+            return try executePaged(cql, pageSize: Self.queryPageSize, consistency: consistency, session: session,
+                                    cancellation: cancellation, startTime: startTime)
         } catch let error as CassandraPluginError where error.refusesPaging {
-            return try executePaged(cql, pageSize: nil, session: session, cancellation: cancellation,
-                                    startTime: startTime)
+            return try executePaged(cql, pageSize: nil, consistency: consistency, session: session,
+                                    cancellation: cancellation, startTime: startTime)
         }
     }
 
     private func executePaged(
         _ cql: String,
         pageSize: Int32?,
+        consistency: CassandraConsistency?,
         session: OpaquePointer,
         cancellation: CassandraCancellation?,
         startTime: Date
@@ -263,6 +268,7 @@
             throw CassandraPluginError.queryFailed("Failed to create statement")
         }
         defer { cass_statement_free(statement) }
+        try Self.apply(consistency, to: statement)
         if let pageSize {
             cass_statement_set_paging_size(statement, pageSize)
         }
@@ -297,6 +303,7 @@
     func executePrepared(
         _ cql: String,
         parameters: [PluginCellValue],
+        consistency: CassandraConsistency?,
         cancellation: CassandraCancellation? = nil
     ) throws -> CassandraRawResult {
         guard let session else {
@@ -316,6 +323,7 @@
         defer { cass_statement_free(statement) }
 
         try CassandraStatementBinder.bind(parameters, to: statement, prepared: prepared)
+        try Self.apply(consistency, to: statement)
         try cancellation?.check()
 
         guard let result = try executePage(statement, on: session, cancellation: cancellation) else {
@@ -336,6 +344,7 @@
     /// is counted by the rows it carried rather than by its size, because ScyllaDB answers a filtered page short.
     func executeBrowse(
         _ browse: CassandraBrowseStatement,
+        consistency: CassandraConsistency?,
         cancellation: CassandraCancellation
     ) throws -> CassandraRawResult {
         if let refusal = browse.window.refusal {
@@ -357,7 +366,7 @@
         if browse.window.offset == 0 {
             resumePoints.remove(key)
         }
-        let walk = CassandraBrowseWalk(browse: browse, pageSize: pageSize, key: key)
+        let walk = CassandraBrowseWalk(browse: browse, pageSize: pageSize, key: key, consistency: consistency)
         let resume = resumePoints.nearest(atOrBefore: browse.window.offset, for: key)
 
         let read: (header: (columns: [String], typeNames: [String])?, rows: [[PluginCellValue]])
@@ -391,6 +400,7 @@
         defer { cass_statement_free(statement) }
 
         try CassandraStatementBinder.bind(walk.browse.window.values.map { .text($0) }, to: statement, prepared: prepared)
+        try Self.apply(walk.consistency, to: statement)
         cass_statement_set_paging_size(statement, Int32(walk.pageSize))
 
         let offset = walk.browse.window.offset
@@ -542,7 +552,7 @@
     }
 
     func switchKeyspace(_ keyspace: String) throws {
-        _ = try executeQuery("USE \"\(escapeIdentifier(keyspace))\"")
+        _ = try executeQuery("USE \"\(escapeIdentifier(keyspace))\"", consistency: nil)
         currentKeyspace = keyspace
     }
 
@@ -944,6 +954,7 @@
     /// rows rather than mid-request.
     func streamQuery(
         _ cql: String,
+        consistency: CassandraConsistency?,
         abort: PluginStreamAbort,
         continuation: AsyncThrowingStream<PluginStreamElement, Error>.Continuation
     ) throws {
@@ -960,12 +971,13 @@
             throw CassandraPluginError.queryFailed("Failed to create statement")
         }
 
+        defer { cass_statement_free(statement) }
+
         cass_statement_set_paging_size(statement, pageSize)
+        try Self.apply(consistency, to: statement)
 
         var headerSent = false
 
-        defer { cass_statement_free(statement) }
-
         var aborted = false
 
         while true {
@@ -1057,6 +1069,17 @@
         value.replacingOccurrences(of: "\"", with: "\"\"")
     }
 
+    private static func apply(_ consistency: CassandraConsistency?, to statement: OpaquePointer) throws {
+        guard let consistency else { return }
+        guard cass_statement_set_consistency(statement, consistency.level.nativeValue) == CASS_OK else {
+            throw CassandraPluginError.queryFailed("Failed to set the consistency level")
+        }
+        guard let serialLevel = consistency.serialLevel else { return }
+        guard cass_statement_set_serial_consistency(statement, serialLevel.nativeValue) == CASS_OK else {
+            throw CassandraPluginError.queryFailed("Failed to set the serial consistency level")
+        }
+    }
+
     static func classifySSLError(rc: CassError, message: String) -> SSLHandshakeError? {
         switch rc {
         case CASS_ERROR_SSL_NO_PEER_CERT, CASS_ERROR_SSL_INVALID_PEER_CERT:
@@ -1087,3 +1110,21 @@
     let rowsAffected: Int
     let executionTime: TimeInterval
 }
+
+extension CassandraConsistencyLevel {
+    var nativeValue: CassConsistency {
+        switch self {
+        case .any: return CASS_CONSISTENCY_ANY
+        case .one: return CASS_CONSISTENCY_ONE
+        case .two: return CASS_CONSISTENCY_TWO
+        case .three: return CASS_CONSISTENCY_THREE
+        case .quorum: return CASS_CONSISTENCY_QUORUM
+        case .all: return CASS_CONSISTENCY_ALL
+        case .localQuorum: return CASS_CONSISTENCY_LOCAL_QUORUM
+        case .eachQuorum: return CASS_CONSISTENCY_EACH_QUORUM
+        case .serial: return CASS_CONSISTENCY_SERIAL
+        case .localSerial: return CASS_CONSISTENCY_LOCAL_SERIAL
+        case .localOne: return CASS_CONSISTENCY_LOCAL_ONE
+        }
+    }
+}

apple/Plugins/CassandraDriverPlugin/CassandraConsistency.swift

@@ -0,0 +1,136 @@
+import Foundation
+import TableProPluginKit
+
+enum CassandraConsistencyLevel: String, CaseIterable, Sendable {
+    case any = "ANY"
+    case one = "ONE"
+    case two = "TWO"
+    case three = "THREE"
+    case quorum = "QUORUM"
+    case all = "ALL"
+    case localQuorum = "LOCAL_QUORUM"
+    case eachQuorum = "EACH_QUORUM"
+    case serial = "SERIAL"
+    case localSerial = "LOCAL_SERIAL"
+    case localOne = "LOCAL_ONE"
+
+    static let serialLevels: [CassandraConsistencyLevel] = [.serial, .localSerial]
+
+    // SERIAL and LOCAL_SERIAL refuse every plain write, so a connection never defaults to them.
+    static let connectionLevels: [CassandraConsistencyLevel] = [
+        .localOne, .one, .two, .three, .quorum, .localQuorum, .eachQuorum, .all, .any,
+    ]
+
+    init?(name: some StringProtocol) {
+        self.init(rawValue: name.uppercased())
+    }
+}
+
+// A nil serial level is never sent, and the server then applies SERIAL.
+struct CassandraConsistency: Sendable, Equatable {
+    static let levelFieldId = "cassandraConsistency"
+    static let serialLevelFieldId = "cassandraSerialConsistency"
+
+    var level: CassandraConsistencyLevel
+    var serialLevel: CassandraConsistencyLevel?
+
+    var effectiveSerialLevel: CassandraConsistencyLevel {
+        serialLevel ?? .serial
+    }
+
+    static var connectionFields: [ConnectionField] {
+        [
+            ConnectionField(
+                id: levelFieldId,
+                label: String(localized: "Consistency"),
+                fieldType: .dropdown(
+                    options: [ConnectionField.DropdownOption(value: "", label: String(localized: "Default"))]
+                        + CassandraConsistencyLevel.connectionLevels.map {
+                            ConnectionField.DropdownOption(value: $0.rawValue, label: $0.rawValue)
+                        }
+                ),
+                section: .advanced
+            ),
+            ConnectionField(
+                id: serialLevelFieldId,
+                label: String(localized: "Serial Consistency"),
+                fieldType: .dropdown(
+                    options: [ConnectionField.DropdownOption(value: "", label: String(localized: "Default"))]
+                        + CassandraConsistencyLevel.serialLevels.map {
+                            ConnectionField.DropdownOption(value: $0.rawValue, label: $0.rawValue)
+                        }
+                ),
+                section: .advanced
+            ),
+        ]
+    }
+
+    init(level: CassandraConsistencyLevel, serialLevel: CassandraConsistencyLevel? = nil) {
+        self.level = level
+        self.serialLevel = serialLevel
+    }
+
+    // Amazon Keyspaces refuses a write at any level but LOCAL_QUORUM, so an unset level there is LOCAL_QUORUM.
+    init(additionalFields: [String: String], host: String) throws {
+        let storedLevel = Self.storedValue(additionalFields[Self.levelFieldId])
+        let storedSerialLevel = Self.storedValue(additionalFields[Self.serialLevelFieldId])
+
+        if let storedLevel {
+            guard let level = CassandraConsistencyLevel(name: storedLevel),
+                  CassandraConsistencyLevel.connectionLevels.contains(level) else {
+                throw CassandraPluginError.connectionFailed(String(
+                    format: String(localized: "%@ is not a consistency level this connection can use. Pick one in the connection's Consistency setting."),
+                    storedLevel
+                ))
+            }
+            self.level = level
+        } else {
+            self.level = Self.isAmazonKeyspaces(host: host, additionalFields: additionalFields) ? .localQuorum : .localOne
+        }
+
+        if let storedSerialLevel {
+            guard let serialLevel = CassandraConsistencyLevel(name: storedSerialLevel),
+                  CassandraConsistencyLevel.serialLevels.contains(serialLevel) else {
+                throw CassandraPluginError.connectionFailed(String(
+                    format: String(localized: "%@ is not a serial consistency level. Use SERIAL or LOCAL_SERIAL."),
+                    storedSerialLevel
+                ))
+            }
+            self.serialLevel = serialLevel
+        } else {
+            self.serialLevel = nil
+        }
+    }
+
+    static func isAmazonKeyspaces(host: String, additionalFields: [String: String]) -> Bool {
+        let awsAuth = additionalFields["awsAuth"] ?? "off"
+        if !awsAuth.isEmpty, awsAuth != "off" { return true }
+        // Through an SSH tunnel the host is the forward's loopback address, and the endpoint stays in preTunnelHost.
+        let endpoint = additionalFields["preTunnelHost"].flatMap { $0.isEmpty ? nil : $0 } ?? host
+        let name = endpoint.trimmingCharacters(in: .whitespaces).lowercased()
+        let trimmed = name.hasSuffix(".") ? String(name.dropLast()) : name
+        return trimmed.range(of: keyspacesEndpointPattern, options: .regularExpression) != nil
+    }
+
+    // cassandra[-fips].<region> on amazonaws.com[.cn] or api.aws, and vpce-<id>.cassandra.<region>.vpce.amazonaws.com.
+    private static let keyspacesEndpointPattern =
+        #"(^|\.)cassandra(-fips)?\.[a-z0-9-]+\.(amazonaws\.com(\.cn)?|api\.aws|vpce\.amazonaws\.com)$"#
+
+    func statusMessage(columns: [String]) -> String {
+        guard columns.first == CassandraRowWriter.unappliedMarker else {
+            return String(format: String(localized: "Consistency: %@"), level.rawValue)
+        }
+        return String(
+            format: String(localized: "Consistency: %1$@, serial: %2$@"),
+            level.rawValue,
+            effectiveSerialLevel.rawValue
+        )
+    }
+
+    private static func storedValue(_ raw: String?) -> String? {
+        guard let trimmed = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
+            return nil
+        }
+        return trimmed
+    }
+}