Decode job from Postgres JSON

Author: adam-fowler

Sources/JobsPostgres/AnyDecodableJob+postgres.swift

@@ -0,0 +1,39 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Hummingbird server framework project
+//
+// Copyright (c) 2024-2025 the Hummingbird authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See hummingbird/CONTRIBUTORS.txt for the list of Hummingbird authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Jobs
+import PostgresNIO
+
+extension AnyDecodableJob: @retroactive PostgresDecodable {
+    public static var psqlType: PostgresDataType {
+        .bytea
+    }
+    public static var psqlFormat: PostgresFormat {
+        .binary
+    }
+
+    public init<JSONDecoder: PostgresJSONDecoder>(
+        from buffer: inout ByteBuffer,
+        type: PostgresDataType,
+        format: PostgresFormat,
+        context: PostgresDecodingContext<JSONDecoder>
+    ) throws {
+        switch (format, type) {
+        case (.binary, .bytea):
+            self = try context.jsonDecoder.decode(Self.self, from: buffer)
+        default:
+            throw PostgresDecodingError.Code.typeMismatch
+        }
+    }
+}

Sources/JobsPostgres/JobRequest+postgres.swift

@@ -0,0 +1,33 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Hummingbird server framework project
+//
+// Copyright (c) 2024-2025 the Hummingbird authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See hummingbird/CONTRIBUTORS.txt for the list of Hummingbird authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Jobs
+import PostgresNIO
+
+extension JobRequest: @retroactive PostgresEncodable where Parameters: Encodable {
+    public static var psqlType: PostgresDataType {
+        .bytea
+    }
+    public static var psqlFormat: PostgresFormat {
+        .binary
+    }
+
+    @inlinable
+    public func encode<JSONEncoder: PostgresJSONEncoder>(
+        into byteBuffer: inout ByteBuffer,
+        context: PostgresEncodingContext<JSONEncoder>
+    ) throws {
+        try context.jsonEncoder.encode(self, into: &byteBuffer)
+    }
+}

Sources/JobsPostgres/PostgresJobsQueue.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Hummingbird server framework project
 //
-// Copyright (c) 2024 the Hummingbird authors
+// Copyright (c) 2024-2025 the Hummingbird authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -173,10 +173,10 @@ public final class PostgresJobQueue: JobQueueDriver {
     /// Push Job onto queue
     /// - Returns: Identifier of queued job
     @discardableResult public func push<Parameters>(_ jobRequest: JobRequest<Parameters>, options: JobOptions) async throws -> JobID {
-        let buffer = try self.jobRegistry.encode(jobRequest: jobRequest)
+        //.let buffer = try self.jobRegistry.encode(jobRequest: jobRequest)
         let jobID = JobID()
         try await self.client.withTransaction(logger: self.logger) { connection in
-            try await self.add(jobID: jobID, jobBuffer: buffer, connection: connection)
+            try await self.add(jobID: jobID, jobRequest: jobRequest, connection: connection)
             try await self.addToQueue(jobID: jobID, connection: connection, delayUntil: options.delayUntil)
         }
         return jobID
@@ -233,10 +233,15 @@ public final class PostgresJobQueue: JobQueueDriver {
     }
 
     func popFirst() async throws -> JobQueueResult<JobID>? {
+        enum PopFirstResult {
+            case nothing
+            case result(Result<PostgresRow, Error>, jobID: JobID)
+        }
         do {
             // The withTransaction closure returns a Result<(ByteBuffer, JobID)?, Error> because
             // we want to be able to exit the closure without cancelling the transaction
-            let result = try await self.client.withTransaction(logger: self.logger) { connection -> Result<(ByteBuffer, JobID)?, Error> in
+            let popFirstResult = try await self.client.withTransaction(logger: self.logger) {
+                connection -> PopFirstResult in
                 try Task.checkCancellation()
 
                 let stream = try await connection.query(
@@ -259,7 +264,7 @@ public final class PostgresJobQueue: JobQueueDriver {
                 )
                 // return nil if nothing in queue
                 guard let jobID = try await stream.decode(UUID.self, context: .default).first(where: { _ in true }) else {
-                    return Result.success(nil)
+                    return .nothing
                 }
                 // set job status to processing
                 try await self.setStatus(jobID: jobID, status: .processing, connection: connection)
@@ -269,25 +274,30 @@ public final class PostgresJobQueue: JobQueueDriver {
                     "SELECT job FROM _hb_pg_jobs WHERE id = \(jobID)",
                     logger: self.logger
                 )
-                guard let buffer = try await stream2.decode(ByteBuffer.self, context: .default).first(where: { _ in true }) else {
+                guard let row = try await stream2.first(where: { _ in true }) else {
                     logger.error(
                         "Failed to find job with id",
                         metadata: [
                             "JobID": "\(jobID)"
                         ]
                     )
-                    // if failed to find the job in the job table return nil
-                    return .success(nil)
+                    // if failed to find the job in the job table return error
+                    return .result(.failure(JobQueueError(code: .unrecognisedJobId, jobName: nil)), jobID: jobID)
                 }
-                return .success((buffer, jobID))
-
+                return .result(.success(row), jobID: jobID)
             }
-            guard let (buffer, jobID) = try result.get() else { return nil }
-            do {
-                let jobInstance = try self.jobRegistry.decode(buffer)
-                return JobQueueResult(id: jobID, result: .success(jobInstance))
-            } catch let error as JobQueueError {
-                return JobQueueResult(id: jobID, result: .failure(error))
+
+            switch popFirstResult {
+            case .nothing:
+                return nil
+            case .result(let result, let jobID):
+                do {
+                    let row = try result.get()
+                    let jobInstance = try row.decode(AnyDecodableJob.self, context: .withJobRegistry(self.jobRegistry)).job
+                    return JobQueueResult(id: jobID, result: .success(jobInstance))
+                } catch let error as JobQueueError {
+                    return JobQueueResult(id: jobID, result: .failure(error))
+                }
             }
         } catch let error as PSQLError {
             logger.error(
@@ -308,11 +318,11 @@ public final class PostgresJobQueue: JobQueueDriver {
         }
     }
 
-    func add(jobID: JobID, jobBuffer: ByteBuffer, connection: PostgresConnection) async throws {
+    func add<Parameters>(jobID: JobID, jobRequest: JobRequest<Parameters>, connection: PostgresConnection) async throws {
         try await connection.query(
             """
             INSERT INTO _hb_pg_jobs (id, job, status)
-            VALUES (\(jobID), \(jobBuffer), \(Status.pending))
+            VALUES (\(jobID), \(jobRequest), \(Status.pending))
             """,
             logger: self.logger
         )
@@ -443,3 +453,12 @@ extension JobQueueDriver where Self == PostgresJobQueue {
         await Self(client: client, migrations: migrations, configuration: configuration, logger: logger)
     }
 }
+
+extension PostgresDecodingContext where JSONDecoder == Foundation.JSONDecoder {
+    /// A ``PostgresDecodingContext`` that uses a Foundation `JSONDecoder` with job registry attached as userInfo.
+    public static func withJobRegistry(_ jobRegistry: JobRegistry) -> PostgresDecodingContext {
+        let jsonDecoder = JSONDecoder()
+        jsonDecoder.userInfo[.configuration] = jobRegistry
+        return PostgresDecodingContext(jsonDecoder: jsonDecoder)
+    }
+}