refactor: migrate to new graphile_worker Plugin API

ceejbot · ceejbot · commit 1eed953b0a53 · 2026-01-23T15:48:57.000-08:00
The graphile_worker_lifecycle_hooks crate changed from a trait-based API
(LifecycleHooks) to an event-registration API (Plugin trait with HookRegistry).

Changes:
- Update lib.rs re-exports to include Plugin, HookRegistry, and event types
- Migrate DlqCleanupPlugin to use hooks.on(JobComplete, ...) pattern
- Update worker.rs add_plugin signature and documentation
- Migrate MetricsPlugin example to new API
- Migrate all test plugins (CountingPlugin, DurationCheckPlugin, etc.)
- Fix hardcoded test database URLs to use consistent defaults
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,7 +16,7 @@ axum = ["dep:axum"]
 
 [dependencies]
 chrono = { version = "0.4", features = ["serde"] }
-graphile_worker = "0.9.2"
+graphile_worker = "0.11.1"
 graphile_worker_crontab_parser = "0.5.12"
 log = "0.4.28"
 metrics = "0.24"
@@ -32,6 +32,7 @@ uuid = { version = "1.0", features = ["v4", "v7", "serde"] }
 
 # Optional dependencies for admin API
 axum = { version = "0.8", optional = true }
+graphile_worker_lifecycle_hooks = "0.2.2"
 
 [dev-dependencies]
 env_logger = "0.11"
diff --git a/examples/metrics_plugin.rs b/examples/metrics_plugin.rs
@@ -20,9 +20,10 @@
 use std::time::Duration;
 
 use backfill::{
-    BackfillClient, BackfillError, IntoTaskHandlerResult, JobCompleteContext, JobFailContext,
-    JobPermanentlyFailContext, JobStartContext, LifecycleHooks, TaskHandler, WorkerConfig, WorkerContext, WorkerRunner,
-    WorkerShutdownContext, WorkerStartContext, enqueue_fast,
+    BackfillClient, BackfillError, HookRegistry, IntoTaskHandlerResult, JobComplete, JobCompleteContext, JobFail,
+    JobFailContext, JobPermanentlyFail, JobPermanentlyFailContext, JobStart, JobStartContext, Plugin, TaskHandler,
+    WorkerConfig, WorkerContext, WorkerRunner, WorkerShutdown, WorkerShutdownContext, WorkerStart, WorkerStartContext,
+    enqueue_fast,
 };
 use serde::{Deserialize, Serialize};
 use tokio_util::sync::CancellationToken;
@@ -37,95 +38,97 @@ use tokio_util::sync::CancellationToken;
 #[derive(Clone)]
 struct MetricsPlugin;
 
-impl LifecycleHooks for MetricsPlugin {
-    async fn on_worker_start(&self, ctx: WorkerStartContext) {
-        metrics::gauge!("backfill_worker_active").set(1.0);
-        log::info!("Worker {} started - metrics recording enabled", ctx.worker_id);
-    }
-
-    async fn on_worker_shutdown(&self, ctx: WorkerShutdownContext) {
-        metrics::gauge!("backfill_worker_active").set(0.0);
-        log::info!("Worker {} shutdown (reason: {:?})", ctx.worker_id, ctx.reason);
-    }
-
-    async fn on_job_start(&self, ctx: JobStartContext) {
-        let task = ctx.job.task_identifier();
-        let attempt = ctx.job.attempts();
-
-        metrics::counter!("jobs_started", "task" => task.clone()).increment(1);
-
-        log::debug!("Job started: {} (attempt {})", task, attempt);
-
-        // Track wait time (time from creation to start)
-        let created_at = ctx.job.created_at();
-        let wait_time = chrono::Utc::now().signed_duration_since(*created_at).num_milliseconds() as f64 / 1000.0;
-
-        metrics::histogram!("job_wait_time_seconds", "task" => task.clone()).record(wait_time);
-    }
-
-    async fn on_job_complete(&self, ctx: JobCompleteContext) {
-        let task = ctx.job.task_identifier();
-        let attempt = ctx.job.attempts();
-        let duration = ctx.duration.as_secs_f64();
-
-        metrics::counter!("jobs_completed", "task" => task.clone(), "attempt" => attempt.to_string()).increment(1);
-
-        metrics::histogram!("job_duration_seconds", "task" => task.clone(), "status" => "success").record(duration);
-
-        log::info!(
-            "Job completed: {} (attempt {}, duration: {:.2}s)",
-            task,
-            attempt,
-            duration
-        );
-    }
-
-    async fn on_job_fail(&self, ctx: JobFailContext) {
-        let task = ctx.job.task_identifier();
-        let attempt = ctx.job.attempts();
-        let will_retry = ctx.will_retry;
-
-        // Use will_retry to distinguish between transient failures and final failures
-        let status = if will_retry { "retrying" } else { "failed" };
-
-        metrics::counter!(
-            "jobs_failed",
-            "task" => task.clone(),
-            "attempt" => attempt.to_string(),
-            "will_retry" => status
-        )
-        .increment(1);
-
-        // Classify the error type for more detailed metrics
-        let error_type = classify_error(&ctx.error);
-        metrics::counter!(
-            "job_errors_by_type",
-            "task" => task.clone(),
-            "error_type" => error_type
-        )
-        .increment(1);
-
-        log::warn!(
-            "Job failed: {} (attempt {}, will_retry: {}, error: {})",
-            task,
-            attempt,
-            will_retry,
-            ctx.error
-        );
-    }
-
-    async fn on_job_permanently_fail(&self, ctx: JobPermanentlyFailContext) {
-        let task = ctx.job.task_identifier();
-        let final_attempt = ctx.job.attempts();
-
-        metrics::counter!("jobs_permanently_failed", "task" => task.clone()).increment(1);
-
-        log::error!(
-            "Job permanently failed: {} (final attempt: {}, error: {})",
-            task,
-            final_attempt,
-            ctx.error
-        );
+impl Plugin for MetricsPlugin {
+    fn register(self, hooks: &mut HookRegistry) {
+        hooks.on(WorkerStart, |ctx: WorkerStartContext| async move {
+            metrics::gauge!("backfill_worker_active").set(1.0);
+            log::info!("Worker {} started - metrics recording enabled", ctx.worker_id);
+        });
+
+        hooks.on(WorkerShutdown, |ctx: WorkerShutdownContext| async move {
+            metrics::gauge!("backfill_worker_active").set(0.0);
+            log::info!("Worker {} shutdown (reason: {:?})", ctx.worker_id, ctx.reason);
+        });
+
+        hooks.on(JobStart, |ctx: JobStartContext| async move {
+            let task = ctx.job.task_identifier();
+            let attempt = ctx.job.attempts();
+
+            metrics::counter!("jobs_started", "task" => task.clone()).increment(1);
+
+            log::debug!("Job started: {} (attempt {})", task, attempt);
+
+            // Track wait time (time from creation to start)
+            let created_at = ctx.job.created_at();
+            let wait_time = chrono::Utc::now().signed_duration_since(*created_at).num_milliseconds() as f64 / 1000.0;
+
+            metrics::histogram!("job_wait_time_seconds", "task" => task.clone()).record(wait_time);
+        });
+
+        hooks.on(JobComplete, |ctx: JobCompleteContext| async move {
+            let task = ctx.job.task_identifier();
+            let attempt = ctx.job.attempts();
+            let duration = ctx.duration.as_secs_f64();
+
+            metrics::counter!("jobs_completed", "task" => task.clone(), "attempt" => attempt.to_string()).increment(1);
+
+            metrics::histogram!("job_duration_seconds", "task" => task.clone(), "status" => "success").record(duration);
+
+            log::info!(
+                "Job completed: {} (attempt {}, duration: {:.2}s)",
+                task,
+                attempt,
+                duration
+            );
+        });
+
+        hooks.on(JobFail, |ctx: JobFailContext| async move {
+            let task = ctx.job.task_identifier();
+            let attempt = ctx.job.attempts();
+            let will_retry = ctx.will_retry;
+
+            // Use will_retry to distinguish between transient failures and final failures
+            let status = if will_retry { "retrying" } else { "failed" };
+
+            metrics::counter!(
+                "jobs_failed",
+                "task" => task.clone(),
+                "attempt" => attempt.to_string(),
+                "will_retry" => status
+            )
+            .increment(1);
+
+            // Classify the error type for more detailed metrics
+            let error_type = classify_error(&ctx.error);
+            metrics::counter!(
+                "job_errors_by_type",
+                "task" => task.clone(),
+                "error_type" => error_type
+            )
+            .increment(1);
+
+            log::warn!(
+                "Job failed: {} (attempt {}, will_retry: {}, error: {})",
+                task,
+                attempt,
+                will_retry,
+                ctx.error
+            );
+        });
+
+        hooks.on(JobPermanentlyFail, |ctx: JobPermanentlyFailContext| async move {
+            let task = ctx.job.task_identifier();
+            let final_attempt = ctx.job.attempts();
+
+            metrics::counter!("jobs_permanently_failed", "task" => task.clone()).increment(1);
+
+            log::error!(
+                "Job permanently failed: {} (final attempt: {}, error: {})",
+                task,
+                final_attempt,
+                ctx.error
+            );
+        });
     }
 }
 
diff --git a/src/dlq_cleanup_plugin.rs b/src/dlq_cleanup_plugin.rs
@@ -8,7 +8,7 @@
 //! The cleanup is based on `job_key` - if a successful job has a `job_key`, any
 //! DLQ entries with that same `job_key` are deleted.
 
-use graphile_worker::{JobCompleteContext, LifecycleHooks};
+use graphile_worker::{HookRegistry, JobComplete, JobCompleteContext, Plugin};
 
 use crate::BackfillClient;
 
@@ -29,31 +29,37 @@ impl DlqCleanupPlugin {
     }
 }
 
-impl LifecycleHooks for DlqCleanupPlugin {
-    async fn on_job_complete(&self, ctx: JobCompleteContext) {
-        // Only clean up if the job has a job_key
-        if let Some(job_key) = ctx.job.key() {
-            match self.client.delete_dlq_by_job_key(job_key).await {
-                Ok(deleted) if deleted > 0 => {
-                    log::debug!(
-                        "Cleaned up DLQ entry after successful job (job_key: {}, task: {})",
-                        job_key,
-                        ctx.job.task_identifier()
-                    );
-                }
-                Ok(_) => {
-                    // No DLQ entry to clean up, this is normal for jobs that
-                    // never failed
-                }
-                Err(e) => {
-                    // Log but don't fail - this is a best-effort cleanup
-                    log::warn!(
-                        "Failed to clean up DLQ entry after job success (job_key: {}, error: {})",
-                        job_key,
-                        e
-                    );
+impl Plugin for DlqCleanupPlugin {
+    fn register(self, hooks: &mut HookRegistry) {
+        let client = self.client;
+        hooks.on(JobComplete, move |ctx: JobCompleteContext| {
+            let client = client.clone();
+            async move {
+                // Only clean up if the job has a job_key
+                if let Some(job_key) = ctx.job.key() {
+                    match client.delete_dlq_by_job_key(job_key).await {
+                        Ok(deleted) if deleted > 0 => {
+                            log::debug!(
+                                "Cleaned up DLQ entry after successful job (job_key: {}, task: {})",
+                                job_key,
+                                ctx.job.task_identifier()
+                            );
+                        }
+                        Ok(_) => {
+                            // No DLQ entry to clean up, this is normal for jobs
+                            // that never failed
+                        }
+                        Err(e) => {
+                            // Log but don't fail - this is a best-effort cleanup
+                            log::warn!(
+                                "Failed to clean up DLQ entry after job success (job_key: {}, error: {})",
+                                job_key,
+                                e
+                            );
+                        }
+                    }
                 }
             }
-        }
+        });
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -60,29 +60,45 @@
 //! ```
 
 use chrono::{DateTime, Utc};
-// Lifecycle hooks for plugins
+// Lifecycle hooks for plugins - new Plugin API with event registration
 pub use graphile_worker::{
+    // Event types (for hooks.on() registration)
+    AfterJobRun,
+    // Context types (for hook handlers)
     AfterJobRunContext,
+    BeforeJobRun,
     BeforeJobRunContext,
+    BeforeJobSchedule,
     BeforeJobScheduleContext,
+    CronJobScheduled,
     CronJobScheduledContext,
+    CronTick,
     CronTickContext,
+    // Plugin trait and registry
+    HookRegistry,
     // Result types
     HookResult,
     // Job type for accessing job data
     Job,
+    JobComplete,
     JobCompleteContext,
+    JobFail,
     JobFailContext,
+    JobFetch,
     JobFetchContext,
+    JobPermanentlyFail,
     JobPermanentlyFailContext,
     JobScheduleResult,
+    JobStart,
     JobStartContext,
-    // Trait
-    LifecycleHooks,
+    Plugin,
+    // Other
     ShutdownReason,
-    // Context types
+    WorkerInit,
     WorkerInitContext,
+    WorkerShutdown,
     WorkerShutdownContext,
+    WorkerStart,
     WorkerStartContext,
 };
 // Re-export commonly used types from graphile_worker
diff --git a/src/worker.rs b/src/worker.rs
@@ -98,7 +98,7 @@ use std::time::Duration;
 use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 
-use crate::{BackfillClient, BackfillError, LifecycleHooks, TaskHandler, WorkerOptions};
+use crate::{BackfillClient, BackfillError, Plugin, TaskHandler, WorkerOptions};
 
 // Re-export for use in wrapper
 /// Configuration for a worker queue
@@ -604,20 +604,28 @@ impl WorkerRunnerBuilder {
     /// Multiple plugins can be registered and will be called in registration
     /// order.
     ///
+    /// # Plugin API
+    /// Plugins implement the `Plugin` trait which has a single `register`
+    /// method that receives a `HookRegistry`. Use `hooks.on(EventType,
+    /// handler)` to register handlers for specific events.
+    ///
     /// # Requirements
     /// The plugin must implement `Clone` to support worker cloning for
-    /// background tasks.
+    /// background tasks. Any state shared between handlers must be wrapped
+    /// in `Arc`.
     ///
     /// # Example
     /// ```rust,no_run
-    /// use backfill::{WorkerRunner, WorkerConfig, LifecycleHooks, JobCompleteContext};
+    /// use backfill::{WorkerRunner, WorkerConfig, Plugin, HookRegistry, JobComplete, JobCompleteContext};
     ///
     /// #[derive(Clone)]
     /// struct MyPlugin;
     ///
-    /// impl LifecycleHooks for MyPlugin {
-    ///     async fn on_job_complete(&self, ctx: JobCompleteContext) {
-    ///         println!("Job {} completed in {:?}", ctx.job.task_identifier(), ctx.duration);
+    /// impl Plugin for MyPlugin {
+    ///     fn register(self, hooks: &mut HookRegistry) {
+    ///         hooks.on(JobComplete, |ctx: JobCompleteContext| async move {
+    ///             println!("Job {} completed in {:?}", ctx.job.task_identifier(), ctx.duration);
+    ///         });
     ///     }
     /// }
     ///
@@ -629,7 +637,7 @@ impl WorkerRunnerBuilder {
     /// # Ok(())
     /// # }
     /// ```
-    pub fn add_plugin<H: LifecycleHooks + Clone + 'static>(mut self, plugin: H) -> Self {
+    pub fn add_plugin<H: Plugin + Clone + 'static>(mut self, plugin: H) -> Self {
         // Store a closure that captures the plugin and applies it to WorkerOptions
         let applier = Arc::new(move |opts: WorkerOptions| opts.add_plugin(plugin.clone()));
         self.plugin_appliers.push(applier);
diff --git a/tests/integration_tests_clean.rs b/tests/integration_tests_clean.rs
@@ -13,7 +13,7 @@ struct TestJob {
 
 /// Get a test database URL from environment or use default
 fn get_test_database_url() -> String {
-    std::env::var("DATABASE_URL").unwrap_or_else(|_| "postgresql://ceej@localhost:5432/backfill_test".to_string())
+    std::env::var("DATABASE_URL").unwrap_or_else(|_| "postgresql://localhost:5432/backfill_test".to_string())
 }
 
 /// One-time setup of test database
diff --git a/tests/plugin_tests.rs b/tests/plugin_tests.rs

Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ struct TestJob {`
`13`	`13`
`14`	`14`	`/// Get a test database URL from environment or use default`
`15`	`15`	`fn get_test_database_url() -> String {`
`16`		`- std::env::var("DATABASE_URL").unwrap_or_else(\|_\| "postgresql://ceej@localhost:5432/backfill_test".to_string())`
	`16`	`+ std::env::var("DATABASE_URL").unwrap_or_else(\|_\| "postgresql://localhost:5432/backfill_test".to_string())`
`17`	`17`	`}`
`18`	`18`
`19`	`19`	`/// One-time setup of test database`