Revert migrations in tests, don't flag migrations as complete in a re…

…vert
hummingbird-project · Mar 6, 2024 · 345f0b6 · 345f0b6
1 parent 3c94c17
commit 345f0b6
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 47 deletions.
diff --git a/...mingbirdJobsPostgres/CreateJobQueue.swift → ...sPostgres/Migrations/CreateJobQueue.swift b/...mingbirdJobsPostgres/CreateJobQueue.swift → ...sPostgres/Migrations/CreateJobQueue.swift
@@ -18,17 +18,6 @@ import Logging
 
 struct CreateJobQueue: HBPostgresMigration {
     func apply(connection: PostgresConnection, logger: Logger) async throws {
-        try await connection.query(
-            """
-            CREATE TABLE IF NOT EXISTS _hb_jobs (
-                id uuid PRIMARY KEY,
-                job json,
-                status smallint,
-                lastModified TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
-            )     
-            """,
-            logger: logger
-        )
         try await connection.query(
             """
             CREATE TABLE IF NOT EXISTS _hb_job_queue (
@@ -48,10 +37,6 @@ struct CreateJobQueue: HBPostgresMigration {
     }
 
     func revert(connection: PostgresConnection, logger: Logger) async throws {
-        try await connection.query(
-            "DROP TABLE _hb_jobs",
-            logger: logger
-        )
         try await connection.query(
             "DROP TABLE _hb_job_queue",
             logger: logger
@@ -61,8 +46,3 @@ struct CreateJobQueue: HBPostgresMigration {
     var name: String { "_Create_JobQueue_Table_" }
     var group: HBMigrationGroup { .jobQueue }
 }
-
-extension HBMigrationGroup {
-    /// JobQueue migration group
-    public static var jobQueue: Self { .init("_hb_jobqueue") }
-}
diff --git a/Sources/HummingbirdJobsPostgres/Migrations/CreateJobs.swift b/Sources/HummingbirdJobsPostgres/Migrations/CreateJobs.swift
@@ -0,0 +1,48 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Hummingbird server framework project
+//
+// Copyright (c) 2024 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 HummingbirdPostgres
+import Logging
+@_spi(ConnectionPool) import PostgresNIO
+
+struct CreateJobs: HBPostgresMigration {
+    func apply(connection: PostgresConnection, logger: Logger) async throws {
+        try await connection.query(
+            """
+            CREATE TABLE IF NOT EXISTS _hb_jobs (
+                id uuid PRIMARY KEY,
+                job bytea,
+                status smallint,
+                lastModified TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
+            )     
+            """,
+            logger: logger
+        )
+    }
+
+    func revert(connection: PostgresConnection, logger: Logger) async throws {
+        try await connection.query(
+            "DROP TABLE _hb_jobs",
+            logger: logger
+        )
+    }
+
+    var name: String { "_Create_Jobs_Table_" }
+    var group: HBMigrationGroup { .jobQueue }
+}
+
+extension HBMigrationGroup {
+    /// JobQueue migration group
+    public static var jobQueue: Self { .init("_hb_jobqueue") }
+}
diff --git a/Sources/HummingbirdJobsPostgres/PostgresJobsQueue.swift b/Sources/HummingbirdJobsPostgres/PostgresJobsQueue.swift
@@ -87,6 +87,7 @@ public final class HBPostgresQueue: HBJobQueueDriver {
         self.logger = logger
         self.isStopped = .init(false)
         self.migrations = migrations
+        await migrations.add(CreateJobs())
         await migrations.add(CreateJobQueue())
     }
 

diff --git a/Sources/HummingbirdPostgres/Migrations.swift b/Sources/HummingbirdPostgres/Migrations.swift
@@ -65,7 +65,7 @@ public actor HBPostgresMigrations {
     ///   - dryRun: Should migrations actually be applied, or should we just report what would be applied and reverted
     @_spi(ConnectionPool)
     public func apply(client: PostgresClient, groups: [HBMigrationGroup] = [], logger: Logger, dryRun: Bool) async throws {
-        try await self.migrate(client: client, migrations: self.migrations, groups: groups, logger: logger, dryRun: dryRun)
+        try await self.migrate(client: client, migrations: self.migrations, groups: groups, logger: logger, completeMigrations: true, dryRun: dryRun)
     }
 
     /// Revery database migrations
@@ -75,14 +75,15 @@ public actor HBPostgresMigrations {
     ///   - dryRun: Should migrations actually be reverted, or should we just report what would be reverted
     @_spi(ConnectionPool)
     public func revert(client: PostgresClient, groups: [HBMigrationGroup] = [], logger: Logger, dryRun: Bool) async throws {
-        try await self.migrate(client: client, migrations: [], groups: groups, logger: logger, dryRun: dryRun)
+        try await self.migrate(client: client, migrations: [], groups: groups, logger: logger, completeMigrations: false, dryRun: dryRun)
     }
 
     private func migrate(
         client: PostgresClient,
         migrations: [HBPostgresMigration],
         groups: [HBMigrationGroup],
         logger: Logger,
+        completeMigrations: Bool,
         dryRun: Bool
     ) async throws {
         switch self.state {
@@ -151,7 +152,9 @@ public actor HBPostgresMigrations {
             self.setFailed(error)
             throw error
         }
-        self.setCompleted()
+        if completeMigrations {
+            self.setCompleted()
+        }
     }
 
     /// Report if the migration process has completed

diff --git a/Tests/HummingbirdPostgresTests/JobsTests.swift b/Tests/HummingbirdPostgresTests/JobsTests.swift
@@ -52,7 +52,7 @@ final class JobsTests: XCTestCase {
             backgroundLogger: logger
         )
         let postgresMigrations = HBPostgresMigrations()
-        return HBJobQueue(
+        return await HBJobQueue(
             HBPostgresQueue(
                 client: postgresClient,
                 migrations: postgresMigrations,
@@ -70,6 +70,7 @@ final class JobsTests: XCTestCase {
     /// shutdown correctly
     @discardableResult public func testJobQueue<T>(
         jobQueue: HBJobQueue<HBPostgresQueue>,
+        revertMigrations: Bool = false,
         test: (HBJobQueue<HBPostgresQueue>) async throws -> T
     ) async throws -> T {
         do {
@@ -85,8 +86,14 @@ final class JobsTests: XCTestCase {
                     try await serviceGroup.run()
                 }
                 do {
-                    try await postgresMigrations.apply(client: postgresClient, groups: [.jobQueue], logger: logger, dryRun: false)
-                    let value = try await test(postgresJobQueue)
+                    let migrations = jobQueue.queue.migrations
+                    let client = jobQueue.queue.client
+                    let logger = jobQueue.queue.logger
+                    if revertMigrations {
+                        try await migrations.revert(client: client, groups: [.jobQueue], logger: logger, dryRun: false)
+                    }
+                    try await migrations.apply(client: client, groups: [.jobQueue], logger: logger, dryRun: false)
+                    let value = try await test(jobQueue)
                     await serviceGroup.triggerGracefulShutdown()
                     return value
                 } catch let error as PSQLError {
@@ -114,7 +121,7 @@ final class JobsTests: XCTestCase {
         test: (HBJobQueue<HBPostgresQueue>) async throws -> T
     ) async throws -> T {
         let jobQueue = try await self.createJobQueue(numWorkers: numWorkers, configuration: configuration)
-        return try await self.testJobQueue(jobQueue: jobQueue, test: test)
+        return try await self.testJobQueue(jobQueue: jobQueue, revertMigrations: true, test: test)
     }
 
     func testBasic() async throws {
@@ -318,20 +325,21 @@ final class JobsTests: XCTestCase {
             backgroundLogger: logger
         )
         let postgresMigrations = HBPostgresMigrations()
-        let jobQueue = HBJobQueue(
+        let jobQueue = await HBJobQueue(
             .postgres(
-                client: postgresClient, 
-                migrations: postgresMigrations, 
+                client: postgresClient,
+                migrations: postgresMigrations,
                 configuration: .init(failedJobsInitialization: .remove, processingJobsInitialization: .remove),
                 logger: logger
             ),
             numWorkers: 2,
             logger: logger
         )
-        let jobQueue2 = HBJobQueue(
+        let postgresMigrations2 = HBPostgresMigrations()
+        let jobQueue2 = await HBJobQueue(
             HBPostgresQueue(
                 client: postgresClient,
-                migrations: postgresMigrations,
+                migrations: postgresMigrations2,
                 configuration: .init(failedJobsInitialization: .remove, processingJobsInitialization: .remove),
                 logger: logger
             ),
@@ -348,22 +356,22 @@ final class JobsTests: XCTestCase {
                     gracefulShutdownSignals: [.sigterm, .sigint],
                     logger: logger
                 )
-                group.addTask {
-                    try await serviceGroup.run()
-                }
-                try await postgresMigrations.apply(client: postgresClient, groups: [.jobQueue], logger: logger, dryRun: false)
-                try await postgresMigrations2.apply(client: postgresClient, groups: [.jobQueue], logger: logger, dryRun: false)
-                do {
-                    for i in 0..<200 {
-                        try await postgresJobQueue.push(id: jobIdentifer, parameters: i)
-                    }
-                    await self.wait(for: [expectation], timeout: 5)
-                    await serviceGroup.triggerGracefulShutdown()
-                } catch {
-                    XCTFail("\(String(reflecting: error))")
-                    await serviceGroup.triggerGracefulShutdown()
-                    throw error
+            )
+            group.addTask {
+                try await serviceGroup.run()
+            }
+            try await postgresMigrations.apply(client: postgresClient, groups: [.jobQueue], logger: logger, dryRun: false)
+            try await postgresMigrations2.apply(client: postgresClient, groups: [.jobQueue], logger: logger, dryRun: false)
+            do {
+                for i in 0..<200 {
+                    try await jobQueue.push(id: jobIdentifer, parameters: i)
                 }
+                await self.wait(for: [expectation], timeout: 5)
+                await serviceGroup.triggerGracefulShutdown()
+            } catch {
+                XCTFail("\(String(reflecting: error))")
+                await serviceGroup.triggerGracefulShutdown()
+                throw error
             }
         }
     }