From ed096ebfc30a01b11ccc4f1375dbdba34eed7e45 Mon Sep 17 00:00:00 2001 From: Jay Herron Date: Wed, 19 Jan 2022 21:23:57 -0700 Subject: [PATCH] Fixes handling of throwing batch load funcs --- Package.resolved | 4 +- README.md | 3 +- Sources/DataLoader/DataLoader.swift | 51 +++++++++++++-------- Tests/DataLoaderTests/DataLoaderTests.swift | 35 ++++++++++++++ 4 files changed, 70 insertions(+), 23 deletions(-) diff --git a/Package.resolved b/Package.resolved index d36443c..55f18b2 100644 --- a/Package.resolved +++ b/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" } } ] diff --git a/README.md b/README.md index 0be1191..c40de24 100644 --- a/README.md +++ b/README.md @@ -251,7 +251,8 @@ struct UserResolver { class UserContext { let database = ... - let userLoader = DataLoader() { [unowned self] keys in + let userLoader = DataLoader() { [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 }! diff --git a/Sources/DataLoader/DataLoader.swift b/Sources/DataLoader/DataLoader.swift index d67e3c4..20b8d73 100644 --- a/Sources/DataLoader/DataLoader.swift +++ b/Sources/DataLoader/DataLoader.swift @@ -28,7 +28,10 @@ final public class DataLoader { private var dispatchScheduled = false private let lock = Lock() - public init(options: DataLoaderOptions = DataLoaderOptions(), batchLoadFunction: @escaping BatchLoadFunction) { + public init( + options: DataLoaderOptions = DataLoaderOptions(), + batchLoadFunction: @escaping BatchLoadFunction + ) { self.options = options self.batchLoadFunction = batchLoadFunction } @@ -37,7 +40,7 @@ final public class DataLoader { public func load(key: Key, on eventLoopGroup: EventLoopGroup) throws -> EventLoopFuture { 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 { 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 { // 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) } } diff --git a/Tests/DataLoaderTests/DataLoaderTests.swift b/Tests/DataLoaderTests/DataLoaderTests.swift index 3f96a1b..bba2927 100644 --- a/Tests/DataLoaderTests/DataLoaderTests.swift +++ b/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( + 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( + options: DataLoaderOptions() + ) { keys in + throw DataLoaderError.typeError(loaderErrorMessage) + } + + XCTAssertThrowsError( + try throwLoaderAutoExecute.load(key: 1, on: eventLoopGroup).wait(), + loaderErrorMessage + ) + } }