Support rate limiting in mock database.

Examples/Reminders/Schema.swift

@@ -288,7 +288,7 @@ func appDatabase() throws -> any DatabaseWriter {
 
     try RemindersList.createTemporaryTrigger(
       after: .delete { _ in
-        Values($createDefaultRemindersList())
+        Select($createDefaultRemindersList())
       } when: { _ in
         !RemindersList.exists()
       }
@@ -365,7 +365,7 @@ func appDatabase() throws -> any DatabaseWriter {
       after: .update {
         $0.status
       } forEachRow: { _, _ in
-        Values($handleReminderStatusUpdate())
+        Select($handleReminderStatusUpdate())
       } when: { _, new in
         new.status.eq(#bind(.completing))
       }

Sources/SQLiteData/CloudKit/Internal/MockCloudDatabase.swift

@@ -6,16 +6,20 @@
 
   @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
@@ -32,8 +36,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(\.continuousClock) var clock
+      self.clock = AnyClock(clock)
+      self.state = LockIsolated(State(rateLimit: rateLimit, rateLimitDuration: rateLimitDuration))
     }
 
     package func set(container: MockCloudContainer) {
@@ -48,6 +59,7 @@
       let accountStatus = container.accountStatus()
       guard accountStatus == .available
       else { throw ckError(forAccountStatus: accountStatus) }
+      try updateRateLimit()
       let record = try state.withValue { state in
         guard let zone = state.storage[recordID.zoneID]
         else { throw CKError(.zoneNotFound) }
@@ -78,13 +90,16 @@
       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) }
+      $_isRateLimitingDisabled.withValue(true) {
+        for id in ids {
+          results[id] = Result { try record(for: id) }
+        }
       }
       return results
     }
@@ -101,6 +116,7 @@
       let accountStatus = container.accountStatus()
       guard accountStatus == .available
       else { throw ckError(forAccountStatus: accountStatus) }
+      try updateRateLimit()
 
       guard (recordsToSave.count + recordIDsToDelete.count) < 200
       else {
@@ -367,6 +383,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 +409,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 +438,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 +453,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/SQLiteData/CloudKit/Internal/Triggers.swift

@@ -198,7 +198,7 @@
           $0.parentRecordType
         )
       } select: {
-        Values(
+        Select(
           #sql("\(new.primaryKey)"),
           T.tableName,
           zoneName ?? defaultZoneName,
@@ -254,7 +254,7 @@
         ifNotExists: true,
         after: .insert { new in
           validate(recordName: new.recordName)
-          Values(
+          Select(
             syncEngine.$didUpdate(
               recordName: new.recordName,
               zoneName: new.zoneName,
@@ -314,7 +314,7 @@
           }
 
           validate(recordName: new.recordName)
-          Values(
+          Select(
             syncEngine.$didUpdate(
               recordName: new.recordName,
               zoneName: new.zoneName,
@@ -337,7 +337,7 @@
         "\(String.sqliteDataCloudKitSchemaName)_after_delete_on_sqlitedata_icloud_metadata",
         ifNotExists: true,
         after: .update(of: \._isDeleted) { _, new in
-          Values(
+          Select(
             syncEngine.$didDelete(
               recordName: new.recordName,
               record: new.lastKnownServerRecord

Sources/SQLiteData/CloudKit/SyncEngine.swift

@@ -1530,7 +1530,7 @@
 
       let unsyncedRecords =
         await withErrorReporting(.sqliteDataCloudKitFailure) {
-          var unsyncedRecordIDs = try await userDatabase.write { db in
+          var unsyncedRecordIDs = try await metadatabase.read { db in
             Set(
               try UnsyncedRecordID.all
                 .fetchAll(db)

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