Query results limit (#20)

* Results limit internal parameter

* Results limit for database API

Targets/Canopy/Sources/CKDatabaseAPI/CKDatabaseAPI.swift

@@ -74,12 +74,14 @@ class CKDatabaseAPI: CKDatabaseAPIType {
   func queryRecords(
     with query: CKQuery,
     in zoneID: CKRecordZone.ID?,
+    resultsLimit: Int?,
     qos: QualityOfService
   ) async -> Result<[CKRecord], CKRecordError> {
     await QueryRecords.with(
       query,
       recordZoneID: zoneID,
       database: database,
+      resultsLimit: resultsLimit,
       qualityOfService: qos
     )
   }

Targets/Canopy/Sources/CKDatabaseAPI/CKDatabaseAPIType+Shorthand.swift

@@ -45,6 +45,10 @@ public extension CKDatabaseAPIType {
   ///   - with: The `CKQuery` specifying your query. See [CloudKit documentation for CKQuery](https://developer.apple.com/documentation/cloudkit/ckquery)
   ///   about how to construct your query.
   ///   - in: The record zone ID to query the records from. or `nil` to use the default zone.
+  ///   - resultsLimit: An optional number of results to return. By default (if this is `nil`),
+  ///   this function returns all matching results, which may be retrieved across several queries. Specify a results limit
+  ///   to set a cap. Note that this should be equal or lower to the CloudKit limit of maximum results per query, which
+  ///   is around a few hundred.
   ///   - qualityOfService: The desired quality of service of the request. Defaults to `.default` if not provided.
   ///
   /// - Returns:
@@ -53,11 +57,13 @@ public extension CKDatabaseAPIType {
   func queryRecords(
     with query: CKQuery,
     in zoneID: CKRecordZone.ID?,
+    resultsLimit: Int? = nil,
     qualityOfService: QualityOfService = .default
   ) async -> Result<[CKRecord], CKRecordError> {
     await queryRecords(
       with: query,
       in: zoneID,
+      resultsLimit: resultsLimit,
       qos: qualityOfService
     )
   }

Targets/Canopy/Sources/CKDatabaseAPI/CKDatabaseAPIType.swift

@@ -15,10 +15,11 @@ public protocol CKDatabaseAPIType {
   typealias PerRecordProgressBlock = (CKRecord, Double) -> Void
   typealias PerRecordIDProgressBlock = (CKRecord.ID, Double) -> Void
 
-  /// See ``CKDatabaseAPIType/queryRecords(with:in:qualityOfService:)`` for preferred way of calling this API.
+  /// See ``CKDatabaseAPIType/queryRecords(with:in:resultsLimit:qualityOfService:)`` for preferred way of calling this API.
   func queryRecords(
     with query: CKQuery,
     in zoneID: CKRecordZone.ID?,
+    resultsLimit: Int?,
     qos: QualityOfService
   ) async -> Result<[CKRecord], CKRecordError>
 

Targets/Canopy/Sources/Features/QueryRecords.swift

@@ -30,6 +30,7 @@ struct QueryRecords {
     recordZoneID: CKRecordZone.ID?,
     database: CKDatabaseType,
     desiredKeys: [CKRecord.FieldKey]? = nil,
+    resultsLimit: Int? = nil,
     qualityOfService: QualityOfService = .default
   ) async -> Result<[CKRecord], CKRecordError> {
     var startingPoint = QueryOperationStartingPoint.query(query)
@@ -41,6 +42,7 @@ struct QueryRecords {
         recordZoneID: recordZoneID,
         database: database,
         desiredKeys: desiredKeys,
+        resultsLimit: resultsLimit,
         qualityOfService: qualityOfService
       )
 
@@ -53,6 +55,11 @@ struct QueryRecords {
         guard !Task.isCancelled else {
           return .failure(.init(from: CKError(CKError.Code.operationCancelled)))
         }
+        // If there was a results limit, just return the result even if there was a cursor
+        if resultsLimit != nil {
+          return .success(records + newRecords)
+        }
+
         startingPoint = QueryOperationStartingPoint.cursor(cursor)
         records += newRecords
       }
@@ -64,6 +71,7 @@ struct QueryRecords {
     recordZoneID: CKRecordZone.ID?,
     database: CKDatabaseType,
     desiredKeys: [CKRecord.FieldKey]? = nil,
+    resultsLimit: Int? = nil,
     qualityOfService: QualityOfService = .userInitiated
   ) async -> QueryOperationResult {
     await withCheckedContinuation { continuation in
@@ -82,6 +90,10 @@ struct QueryRecords {
       operation.desiredKeys = desiredKeys
       operation.qualityOfService = qualityOfService
 
+      if let resultsLimit {
+        operation.resultsLimit = resultsLimit
+      }
+
       operation.recordMatchedBlock = { _, result in
         switch result {
         case let .failure(error):

Targets/Canopy/Tests/QueryRecordsFeatureTests.swift

@@ -72,6 +72,34 @@ final class QueryRecordsFeatureTests: XCTestCase {
     XCTAssertEqual(records[19].recordID.recordName, "id20")
   }
 
+  func test_results_limit_query() async {
+    let query = CKQuery(recordType: "TestRecord", predicate: NSPredicate(value: true))
+    let db = ReplayingMockCKDatabase(
+      operationResults: [
+        .query(
+          .init(
+            queryRecordResults: records(startIndex: 1, endIndex: 10).map {
+              ReplayingMockCKDatabase.QueryRecordResult(recordID: $0.recordID, result: .success($0))
+            },
+            queryResult: .init(result: .success(CKQueryOperation.Cursor.mock))
+          )
+        )
+      ]
+    )
+
+    let records = try! await QueryRecords.with(
+      query,
+      recordZoneID: nil,
+      database: db,
+      resultsLimit: 10
+    ).get()
+    XCTAssertEqual(records.count, 10)
+    let operationsRun = await db.operationsRun
+    XCTAssertEqual(operationsRun, 1)
+    XCTAssertEqual(records[0].recordID.recordName, "id1")
+    XCTAssertEqual(records[9].recordID.recordName, "id10")
+  }
+
   func test_depth3_query() async {
     let query = CKQuery(recordType: "TestRecord", predicate: NSPredicate(value: true))
     let db = ReplayingMockCKDatabase(