Support rate limiting in mock database.

Package.swift

@@ -61,6 +61,7 @@ let package = Package(
   dependencies: [
     .package(url: "https://github.com/apple/swift-collections", from: "1.0.0"),
     .package(url: "https://github.com/groue/GRDB.swift", from: "7.6.0"),
+    .package(url: "https://github.com/pointfreeco/swift-clocks", from: "1.0.0"),
     .package(url: "https://github.com/pointfreeco/swift-concurrency-extras", from: "1.4.0"),
     .package(url: "https://github.com/pointfreeco/swift-custom-dump", from: "1.3.3"),
     .package(url: "https://github.com/pointfreeco/swift-dependencies", from: "1.9.0"),
@@ -91,6 +92,7 @@ let package = Package(
     .target(
       name: "SQLiteData",
       dependencies: [
+        .product(name: "Clocks", package: "swift-clocks"),
         .product(name: "ConcurrencyExtras", package: "swift-concurrency-extras"),
         .product(name: "Dependencies", package: "swift-dependencies"),
         .product(name: "GRDB", package: "GRDB.swift"),

Package@swift-6.0.swift

@@ -23,6 +23,7 @@ let package = Package(
   dependencies: [
     .package(url: "https://github.com/apple/swift-collections", from: "1.0.0"),
     .package(url: "https://github.com/groue/GRDB.swift", from: "7.6.0"),
+    .package(url: "https://github.com/pointfreeco/swift-clocks", from: "1.0.0"),
     .package(url: "https://github.com/pointfreeco/swift-concurrency-extras", from: "1.0.0"),
     .package(url: "https://github.com/pointfreeco/swift-custom-dump", from: "1.3.3"),
     .package(url: "https://github.com/pointfreeco/swift-dependencies", from: "1.9.0"),
@@ -35,6 +36,7 @@ let package = Package(
     .target(
       name: "SQLiteData",
       dependencies: [
+        .product(name: "Clocks", package: "swift-clocks"),
         .product(name: "ConcurrencyExtras", package: "swift-concurrency-extras"),
         .product(name: "Dependencies", package: "swift-dependencies"),
         .product(name: "GRDB", package: "GRDB.swift"),

Package@swift-6.1.swift

@@ -61,6 +61,7 @@ let package = Package(
   dependencies: [
     .package(url: "https://github.com/apple/swift-collections", from: "1.0.0"),
     .package(url: "https://github.com/groue/GRDB.swift", from: "7.6.0"),
+    .package(url: "https://github.com/pointfreeco/swift-clocks", from: "1.0.0"),
     .package(url: "https://github.com/pointfreeco/swift-concurrency-extras", from: "1.4.0"),
     .package(url: "https://github.com/pointfreeco/swift-custom-dump", from: "1.3.3"),
     .package(url: "https://github.com/pointfreeco/swift-dependencies", from: "1.9.0"),
@@ -82,6 +83,7 @@ let package = Package(
     .target(
       name: "SQLiteData",
       dependencies: [
+        .product(name: "Clocks", package: "swift-clocks"),
         .product(name: "ConcurrencyExtras", package: "swift-concurrency-extras"),
         .product(name: "Dependencies", package: "swift-dependencies"),
         .product(name: "GRDB", package: "GRDB.swift"),

Sources/SQLiteData/CloudKit/Internal/CloudKitClock.swift

@@ -0,0 +1,15 @@
+import Dependencies
+import Clocks
+
+@available(iOS 16, macOS 13, tvOS 16, watchOS 9, *)
+extension DependencyValues {
+  package var cloudKitClock: any Clock<Duration> {
+    get { self[CloudKitClockKey.self] }
+    set { self[CloudKitClockKey.self] = newValue }
+  }
+
+  private enum CloudKitClockKey: DependencyKey {
+    static let liveValue: any Clock<Duration> = ContinuousClock()
+    static let testValue: any Clock<Duration> = ContinuousClock()
+  }
+}

Sources/SQLiteData/CloudKit/Internal/MockCloudDatabase.swift

@@ -6,20 +6,31 @@
 
   @available(iOS 17, macOS 14, tvOS 17, watchOS 10, *)
   package final class MockCloudDatabase: CloudDatabase {
-    package let state = LockIsolated(State())
+    package let state: LockIsolated<State>
     package let databaseScope: CKDatabase.Scope
     let _container = IsolatedWeakVar<MockCloudContainer>()
     let dataManager = Dependency(\.dataManager)
+    let clock: AnyClock<Duration>
 
     package struct State {
       private var lastRecordChangeTag = 0
       package var storage: [CKRecordZone.ID: Zone] = [:]
       var assets: [AssetID: Data] = [:]
       var deletedRecords: [(CKRecord.ID, CKRecord.RecordType)] = []
+      package var rateLimit: Int
+      package var rateLimitDuration: Duration
+      fileprivate var requestTimestamps: [AnyClock<Duration>.Instant] = []
       mutating func nextRecordChangeTag() -> Int {
         lastRecordChangeTag += 1
         return lastRecordChangeTag
       }
+      init(
+        rateLimit: Int,
+        rateLimitDuration: Duration
+      ) {
+        self.rateLimit = rateLimit
+        self.rateLimitDuration = rateLimitDuration
+      }
     }
 
     struct AssetID: Hashable {
@@ -32,8 +43,15 @@
       package var records: [CKRecord.ID: CKRecord] = [:]
     }
 
-    package init(databaseScope: CKDatabase.Scope) {
+    package init(
+      databaseScope: CKDatabase.Scope,
+      rateLimit: Int = .max,
+      rateLimitDuration: Duration = .seconds(1)
+    ) {
       self.databaseScope = databaseScope
+      @Dependency(\.cloudKitClock) var clock
+      self.clock = AnyClock(clock)
+      self.state = LockIsolated(State(rateLimit: rateLimit, rateLimitDuration: rateLimitDuration))
     }
 
     package func set(container: MockCloudContainer) {
@@ -45,29 +63,10 @@
     }
 
     package func record(for recordID: CKRecord.ID) throws -> CKRecord {
-      let accountStatus = container.accountStatus()
-      guard accountStatus == .available
-      else { throw ckError(forAccountStatus: accountStatus) }
-      let record = try state.withValue { state in
-        guard let zone = state.storage[recordID.zoneID]
-        else { throw CKError(.zoneNotFound) }
-        guard let record = zone.records[recordID]
-        else { throw CKError(.unknownItem) }
-        guard let record = record.copy() as? CKRecord
-        else { fatalError("Could not copy CKRecord.") }
-        return record
-      }
-
-      try state.withValue { state in
-        for key in record.allKeys() {
-          guard let assetData = state.assets[AssetID(recordID: record.recordID, key: key)]
-          else { continue }
-          let url = dataManager.wrappedValue.temporaryDirectory.appending(path: UUID().uuidString)
-          try dataManager.wrappedValue.save(assetData, to: url)
-          record[key] = CKAsset(fileURL: url)
-        }
+      guard let record = try records(for: [recordID], desiredKeys: nil).values.first?.get()
+      else {
+        throw CKError(.unknownItem)
       }
-
       return record
     }
 
@@ -78,15 +77,34 @@
       let accountStatus = container.accountStatus()
       guard accountStatus == .available
       else { throw ckError(forAccountStatus: accountStatus) }
-
+      try updateRateLimit()
       guard ids.count < 200
       else { throw CKError(.limitExceeded) }
 
-      var results: [CKRecord.ID: Result<CKRecord, any Error>] = [:]
-      for id in ids {
-        results[id] = Result { try record(for: id) }
+      return state.withValue { state in
+        var result: [CKRecord.ID: Result<CKRecord, any Error>] = [:]
+        for recordID in ids {
+          result[recordID] = Result {
+            guard let zone = state.storage[recordID.zoneID]
+            else { throw CKError(.zoneNotFound) }
+            guard let record = zone.records[recordID]
+            else { throw CKError(.unknownItem) }
+            guard let record = record.copy() as? CKRecord
+            else { fatalError("Could not copy CKRecord.") }
+
+            for key in record.allKeys() {
+              guard let assetData = state.assets[AssetID(recordID: record.recordID, key: key)]
+              else { continue }
+              let url = dataManager.wrappedValue.temporaryDirectory.appending(path: UUID().uuidString)
+              try dataManager.wrappedValue.save(assetData, to: url)
+              record[key] = CKAsset(fileURL: url)
+            }
+
+            return record
+          }
+        }
+        return result
       }
-      return results
     }
 
     package func modifyRecords(
@@ -101,6 +119,7 @@
       let accountStatus = container.accountStatus()
       guard accountStatus == .available
       else { throw ckError(forAccountStatus: accountStatus) }
+      try updateRateLimit()
 
       guard (recordsToSave.count + recordIDsToDelete.count) < 200
       else {
@@ -367,6 +386,7 @@
       let accountStatus = container.accountStatus()
       guard accountStatus == .available
       else { throw ckError(forAccountStatus: accountStatus) }
+      try updateRateLimit()
 
       return state.withValue { state in
         var saveResults: [CKRecordZone.ID: Result<CKRecordZone, any Error>] = [:]
@@ -392,6 +412,26 @@
       }
     }
 
+    private func updateRateLimit() throws {
+      guard !_isRateLimitingDisabled else { return }
+      try state.withValue { state in
+        guard state.rateLimit != .max
+        else { return }
+        let now = clock.now
+        state.requestTimestamps.removeAll { $0.duration(to: now) >= state.rateLimitDuration }
+        let isRateLimited = state.requestTimestamps.count >= state.rateLimit
+        state.requestTimestamps.append(now)
+        guard !isRateLimited
+        else {
+          let (seconds, _) = state.rateLimitDuration.components
+          throw CKError(
+            .requestRateLimited,
+            userInfo: [CKErrorRetryAfterKey: Double(seconds)]
+          )
+        }
+      }
+    }
+
     package nonisolated static func == (lhs: MockCloudDatabase, rhs: MockCloudDatabase) -> Bool {
       lhs === rhs
     }
@@ -401,6 +441,8 @@
     }
   }
 
+  @TaskLocal package var _isRateLimitingDisabled = false
+
   @available(macOS 13, iOS 16, tvOS 16, watchOS 9, *)
   private func ckError(forAccountStatus accountStatus: CKAccountStatus) -> CKError {
     switch accountStatus {
@@ -414,4 +456,5 @@
       fatalError()
     }
   }
+
 #endif

Sources/SQLiteData/CloudKit/Internal/MockSyncEngine.swift

@@ -267,12 +267,14 @@
         return SendRecordsCallback {}
       }
 
-      let (saveResults, deleteResults) = try syncEngine.database.modifyRecords(
-        saving: batch.recordsToSave,
-        deleting: batch.recordIDsToDelete,
-        savePolicy: .ifServerRecordUnchanged,
-        atomically: batch.atomicByZone
-      )
+      let (saveResults, deleteResults) = try $_isRateLimitingDisabled.withValue(true) {
+        try syncEngine.database.modifyRecords(
+          saving: batch.recordsToSave,
+          deleting: batch.recordIDsToDelete,
+          savePolicy: .ifServerRecordUnchanged,
+          atomically: batch.atomicByZone
+        )
+      }
 
       var savedRecords: [CKRecord] = []
       var failedRecordSaves: [(record: CKRecord, error: CKError)] = []
@@ -395,10 +397,12 @@
         (
           saveResults: [CKRecordZone.ID: Result<CKRecordZone, any Error>],
           deleteResults: [CKRecordZone.ID: Result<Void, any Error>]
-        ) = try syncEngine.database.modifyRecordZones(
-          saving: zonesToSave,
-          deleting: zoneIDsToDelete
-        )
+        ) = try $_isRateLimitingDisabled.withValue(true) {
+          try syncEngine.database.modifyRecordZones(
+            saving: zonesToSave,
+            deleting: zoneIDsToDelete
+          )
+        }
       var savedZones: [CKRecordZone] = []
       var failedZoneSaves: [(zone: CKRecordZone, error: CKError)] = []
       var deletedZoneIDs: [CKRecordZone.ID] = []

Sources/TestLocals/TestLocals.swift

@@ -7,6 +7,8 @@
   @TaskLocal package var startImmediately = true
   @TaskLocal package var attachMetadatabase = false
   @TaskLocal package var accountStatus = CKAccountStatus.available
+  @TaskLocal package var rateLimit = Int.max
+  @TaskLocal package var rateLimitDuration = Duration.seconds(1)
   @available(iOS 17, macOS 14, tvOS 17, watchOS 10, *)
   @TaskLocal package var syncEngineDelegate: (any SyncEngineDelegate)? = nil
 #endif