feat: cyclotron fetch and janitor run migrations (#24559)

Co-authored-by: Brett Hoerner <brett@bretthoerner.com>

.github/workflows/rust-cyclotron-migrator-docker.yml

@@ -1,81 +0,0 @@
-name: Build rust cyclotron-migrator docker image
-
-on:
-    workflow_dispatch:
-    push:
-        paths:
-            - 'rust/**'
-            - '.github/workflows/rust-cyclotron-migrator-docker.yml'
-        branches:
-            - 'master'
-
-permissions:
-    packages: write
-
-jobs:
-    build:
-        name: build and publish cyclotron-migrator image
-        runs-on: depot-ubuntu-22.04-4
-        permissions:
-            id-token: write # allow issuing OIDC tokens for this workflow run
-            contents: read # allow reading the repo contents
-            packages: write # allow push to ghcr.io
-
-        defaults:
-            run:
-                working-directory: rust
-
-        steps:
-            - name: Check Out Repo
-              # Checkout project code
-              # Use sparse checkout to only select files in rust directory
-              # Turning off cone mode ensures that files in the project root are not included during checkout
-              uses: actions/checkout@v3
-              with:
-                  sparse-checkout: 'rust/'
-                  sparse-checkout-cone-mode: false
-
-            - name: Set up Depot CLI
-              uses: depot/setup-action@v1
-
-            - name: Login to ghcr.io
-              uses: docker/login-action@v2
-              with:
-                  registry: ghcr.io
-                  username: ${{ github.actor }}
-                  password: ${{ secrets.GITHUB_TOKEN }}
-
-            - name: Set up QEMU
-              uses: docker/setup-qemu-action@v3
-
-            - name: Docker meta
-              id: meta
-              uses: docker/metadata-action@v5
-              with:
-                  images: ghcr.io/posthog/posthog/cyclotron-migrator
-                  tags: |
-                      type=ref,event=pr
-                      type=ref,event=branch
-                      type=semver,pattern={{version}}
-                      type=semver,pattern={{major}}.{{minor}}
-                      type=sha
-
-            - name: Set up Docker Buildx
-              id: buildx
-              uses: docker/setup-buildx-action@v2
-
-            - name: Build and push migrator
-              id: docker_build_cyclotron_migrator
-              uses: depot/build-push-action@v1
-              with:
-                  context: ./rust/
-                  file: ./rust/Dockerfile.migrate-cyclotron
-                  push: true
-                  tags: ${{ steps.meta.outputs.tags }}
-                  labels: ${{ steps.meta.outputs.labels }}
-                  platforms: linux/arm64
-                  cache-from: type=gha
-                  cache-to: type=gha,mode=max
-
-            - name: Cyclotron-migrator image digest
-              run: echo ${{ steps.docker_build_cyclotron_migrator.outputs.digest }}

bin/start-cyclotron

@@ -14,9 +14,6 @@ export RUST_LOG=$RUST_LOG,sqlx::query=$SQLX_QUERY_LEVEL
 
 export DATABASE_URL=${DATABASE_URL:-postgres://posthog:posthog@localhost:5432/posthog}
 export ALLOW_INTERNAL_IPS=${ALLOW_INTERNAL_IPS:-true}
-cd cyclotron-core
-cargo sqlx migrate run
-cd ..
 
 ./target/debug/cyclotron-fetch &
 ./target/debug/cyclotron-janitor &

rust/Dockerfile.migrate-cyclotron

@@ -1,16 +0,0 @@
-FROM docker.io/library/rust:1.80.1-bullseye as builder
-
-RUN apt update && apt install build-essential cmake -y
-RUN cargo install sqlx-cli@0.7.3 --no-default-features --features native-tls,postgres --root /app/target/release/
-
-FROM debian:bullseye-20230320-slim AS runtime
-WORKDIR /sqlx
-
-COPY bin /sqlx/bin/
-COPY cyclotron-core/migrations /sqlx/migrations/
-
-COPY --from=builder /app/target/release/bin/sqlx /usr/local/bin
-
-RUN chmod +x ./bin/migrate
-
-CMD ["./bin/migrate"]

rust/cyclotron-core/src/lib.rs

@@ -36,6 +36,9 @@ mod config;
 pub use config::ManagerConfig;
 pub use config::PoolConfig;
 
+// Meta
+pub use ops::meta::run_migrations;
+
 // Some data is shared between workers and janitors on a given shard, using
 // the metadata table. These keys are reserved for that purpose
 

rust/cyclotron-core/src/ops/meta.rs

@@ -1,4 +1,4 @@
-use sqlx::postgres::PgQueryResult;
+use sqlx::{postgres::PgQueryResult, PgPool};
 use uuid::Uuid;
 
 use crate::error::QueueError;
@@ -24,3 +24,11 @@ pub fn throw_if_no_rows(res: PgQueryResult, job: Uuid, lock: Uuid) -> Result<(),
         Ok(())
     }
 }
+
+/// Run the latest cyclotron migrations. Panics if the migrations can't be run - failure to run migrations is purposefully fatal.
+pub async fn run_migrations(pool: &PgPool) {
+    sqlx::migrate!("./migrations")
+        .run(pool)
+        .await
+        .expect("Failed to run migrations");
+}

rust/cyclotron-core/src/worker.rs

@@ -6,7 +6,10 @@ use std::sync::Mutex;
 use uuid::Uuid;
 
 use crate::{
-    ops::worker::{dequeue_jobs, dequeue_with_vm_state, flush_job, get_vm_state, set_heartbeat},
+    ops::{
+        meta::run_migrations,
+        worker::{dequeue_jobs, dequeue_with_vm_state, flush_job, get_vm_state, set_heartbeat},
+    },
     Job, JobState, JobUpdate, PoolConfig, QueueError,
 };
 
@@ -49,6 +52,11 @@ impl Worker {
         }
     }
 
+    /// Run the latest cyclotron migrations. Panics if the migrations can't be run - failure to run migrations is purposefully fatal.
+    pub async fn run_migrations(&self) {
+        run_migrations(&self.pool).await;
+    }
+
     /// Dequeues jobs from the queue, and returns them. Job sorting happens at the queue level,
     /// workers can't provide any filtering or sorting criteria - queue managers decide which jobs are run,
     /// workers just run them.

rust/cyclotron-fetch/src/main.rs

@@ -75,6 +75,8 @@ async fn main() {
         .await
         .expect("failed to create app context");
 
+    context.worker.run_migrations().await;
+
     let http_server = tokio::spawn(listen(app, bind));
 
     let worker_loop = tokio::spawn(worker_loop(context));

rust/cyclotron-fetch/tests/fetch.rs

@@ -13,6 +13,14 @@ use utils::{
 
 mod utils;
 
+#[sqlx::test(migrations = "../cyclotron-core/migrations")]
+pub async fn test_run_migrations(db: PgPool) {
+    // This is a no-op, since the db sqlx::test gives use already has the migrations run, but it asserts that the migrations
+    // being run repeatedly doesn't cause any issues, and that the migrations being run are the same as the ones in the core
+    let context = get_app_test_context(db).await;
+    context.worker.run_migrations().await;
+}
+
 #[sqlx::test(migrations = "../cyclotron-core/migrations")]
 pub async fn test_completes_fetch(db: PgPool) {
     let context = Arc::new(get_app_test_context(db.clone()).await);

rust/cyclotron-janitor/src/janitor.rs

@@ -56,6 +56,10 @@ impl Janitor {
         }
     }
 
+    pub async fn run_migrations(&self) {
+        cyclotron_core::run_migrations(&self.pool).await;
+    }
+
     pub async fn run_once(&self) -> Result<CleanupResult, QueueError> {
         info!("Running janitor loop");
         let _loop_start = common_metrics::timing_guard(RUN_TIME, &self.metrics_labels);

rust/cyclotron-janitor/src/main.rs

@@ -70,6 +70,8 @@ async fn main() {
         .await
         .expect("failed to create janitor");
 
+    janitor.run_migrations().await;
+
     let janitor_liveness = liveness
         .register(
             "janitor".to_string(),