Merge pull request #13 from NeedleInAJayStack/fix/throwingBatchFunctions

Fixes handling of throwing batch load funcs

Author: NeedleInAJayStack

Package.resolved

@@ -6,8 +6,8 @@
         "repositoryURL": "https://github.com/apple/swift-nio.git",
         "state": {
           "branch": null,
-          "revision": "b8368b6e09b7993896c42a6199103a73ecc1dbf9",
-          "version": "2.0.0"
+          "revision": "51c3fc2e4a0fcdf4a25089b288dd65b73df1b0ef",
+          "version": "2.37.0"
         }
       }
     ]

README.md

@@ -251,7 +251,8 @@ struct UserResolver {
 
 class UserContext {
     let database = ...
-    let userLoader = DataLoader<Int, User>() { [unowned self] keys in
+    let userLoader = DataLoader<Int, User>() { [weak self] keys in
+        guard let self = self else { throw ContextError }
         return User.query(on: self.database).filter(\.$id ~~ keys).all().map { users in
             keys.map { key in
                 users.first { $0.id == key }!

Sources/DataLoader/DataLoader.swift

@@ -28,7 +28,10 @@ final public class DataLoader<Key: Hashable, Value> {
     private var dispatchScheduled = false
     private let lock = Lock()
 
-    public init(options: DataLoaderOptions<Key, Value> = DataLoaderOptions(), batchLoadFunction: @escaping BatchLoadFunction<Key, Value>) {
+    public init(
+        options: DataLoaderOptions<Key, Value> = DataLoaderOptions(),
+        batchLoadFunction: @escaping BatchLoadFunction<Key, Value>
+    ) {
         self.options = options
         self.batchLoadFunction = batchLoadFunction
     }
@@ -37,7 +40,7 @@ final public class DataLoader<Key: Hashable, Value> {
     public func load(key: Key, on eventLoopGroup: EventLoopGroup) throws -> EventLoopFuture<Value> {
         let cacheKey = options.cacheKeyFunction?(key) ?? key
 
-        return try lock.withLock {
+        return lock.withLock {
             if options.cachingEnabled, let cachedFuture = cache[cacheKey] {
                 return cachedFuture
             }
@@ -53,16 +56,20 @@ final public class DataLoader<Key: Hashable, Value> {
                     dispatchScheduled = true
                 }
             } else {
-                _ = try batchLoadFunction([key]).map { results  in
-                    if results.isEmpty {
-                        promise.fail(DataLoaderError.noValueForKey("Did not return value for key: \(key)"))
-                    } else {
-                        let result = results[0]
-                        switch result {
-                        case .success(let value): promise.succeed(value)
-                        case .failure(let error): promise.fail(error)
+                do {
+                    _ = try batchLoadFunction([key]).map { results  in
+                        if results.isEmpty {
+                            promise.fail(DataLoaderError.noValueForKey("Did not return value for key: \(key)"))
+                        } else {
+                            let result = results[0]
+                            switch result {
+                            case .success(let value): promise.succeed(value)
+                            case .failure(let error): promise.fail(error)
+                            }
                         }
                     }
+                } catch {
+                    promise.fail(error)
                 }
             }
 
@@ -181,20 +188,24 @@ final public class DataLoader<Key: Hashable, Value> {
 
         // Step through the values, resolving or rejecting each Promise in the
         // loaded queue.
-        _ = try batchLoadFunction(keys).flatMapThrowing { values in
-            if values.count != keys.count {
-                throw DataLoaderError.typeError("The function did not return an array of the same length as the array of keys. \nKeys count: \(keys.count)\nValues count: \(values.count)")
-            }
+        do {
+            _ = try batchLoadFunction(keys).flatMapThrowing { values in
+                if values.count != keys.count {
+                    throw DataLoaderError.typeError("The function did not return an array of the same length as the array of keys. \nKeys count: \(keys.count)\nValues count: \(values.count)")
+                }
 
-            for entry in batch.enumerated() {
-                let result = values[entry.offset]
+                for entry in batch.enumerated() {
+                    let result = values[entry.offset]
 
-                switch result {
-                case .failure(let error): entry.element.promise.fail(error)
-                case .success(let value): entry.element.promise.succeed(value)
+                    switch result {
+                    case .failure(let error): entry.element.promise.fail(error)
+                    case .success(let value): entry.element.promise.succeed(value)
+                    }
                 }
+            }.recover { error in
+                self.failedExecution(batch: batch, error: error)
             }
-        }.recover { error in
+        } catch {
             self.failedExecution(batch: batch, error: error)
         }
     }

Tests/DataLoaderTests/DataLoaderTests.swift

@@ -441,4 +441,39 @@ final class DataLoaderTests: XCTestCase {
 
         XCTAssertNotNil(value)
     }
+    
+    func testErrorResult() throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer {
+            XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
+        }
+        
+        let loaderErrorMessage = "TEST"
+        
+        // Test throwing loader without auto-executing
+        let throwLoader = DataLoader<Int, Int>(
+            options: DataLoaderOptions(executionPeriod: nil)
+        ) { keys in
+            throw DataLoaderError.typeError(loaderErrorMessage)
+        }
+        
+        let value = try throwLoader.load(key: 1, on: eventLoopGroup)
+        XCTAssertNoThrow(try throwLoader.execute())
+        XCTAssertThrowsError(
+            try value.wait(),
+            loaderErrorMessage
+        )
+        
+        // Test throwing loader with auto-executing
+        let throwLoaderAutoExecute = DataLoader<Int, Int>(
+            options: DataLoaderOptions()
+        ) { keys in
+            throw DataLoaderError.typeError(loaderErrorMessage)
+        }
+        
+        XCTAssertThrowsError(
+            try throwLoaderAutoExecute.load(key: 1, on: eventLoopGroup).wait(),
+            loaderErrorMessage
+        )
+    }
 }