crates-admin: Add backfill tool for generating OG images for existing crates

Author: Turbo87

src/bin/crates-admin/backfill_og_images.rs

@@ -0,0 +1,113 @@
+use anyhow::Result;
+use crates_io::db;
+use crates_io::schema::crates;
+use crates_io::worker::jobs::GenerateOgImage;
+use crates_io_worker::BackgroundJob;
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use tracing::{info, warn};
+
+#[derive(clap::Parser, Debug)]
+#[command(
+    name = "backfill-og-images",
+    about = "Enqueue OG image generation jobs for existing crates"
+)]
+pub struct Opts {
+    #[arg(long, default_value = "1000")]
+    /// Batch size for enqueueing crates (default: 1000)
+    batch_size: usize,
+
+    #[arg(long)]
+    /// Only generate OG images for crates with names starting with this prefix
+    prefix: Option<String>,
+
+    #[arg(long)]
+    /// Offset to start enqueueing from (useful for resuming)
+    offset: Option<i64>,
+}
+
+pub async fn run(opts: Opts) -> Result<()> {
+    let mut conn = db::oneoff_connection().await?;
+
+    info!("Starting OG image backfill with options: {opts:?}");
+
+    // Helper function to build query
+    let build_query = |offset: i64| {
+        let mut query = crates::table
+            .select(crates::name)
+            .order(crates::name)
+            .into_boxed();
+
+        if let Some(prefix) = &opts.prefix {
+            query = query.filter(crates::name.like(format!("{prefix}%")));
+        }
+
+        query.offset(offset)
+    };
+
+    // Count total crates to process
+    let mut count_query = crates::table.into_boxed();
+    if let Some(prefix) = &opts.prefix {
+        count_query = count_query.filter(crates::name.like(format!("{prefix}%")));
+    }
+    let total_crates: i64 = count_query.count().get_result(&mut conn).await?;
+
+    info!("Total crates to enqueue: {total_crates}");
+
+    let mut offset = opts.offset.unwrap_or(0);
+    let mut enqueued = 0;
+    let mut errors = 0;
+
+    loop {
+        // Fetch batch of crate names
+        let crate_names: Vec<String> = build_query(offset)
+            .limit(opts.batch_size as i64)
+            .load(&mut conn)
+            .await?;
+
+        if crate_names.is_empty() {
+            break;
+        }
+
+        let batch_size = crate_names.len();
+        info!(
+            "Enqueueing batch {}-{} of {total_crates}",
+            offset + 1,
+            offset + batch_size as i64
+        );
+
+        for crate_name in crate_names {
+            match GenerateOgImage::new(crate_name.clone())
+                .enqueue(&mut conn)
+                .await
+            {
+                Ok(_) => {
+                    enqueued += 1;
+                    if enqueued % 100 == 0 {
+                        info!("Enqueued {enqueued} jobs so far...");
+                    }
+                }
+                Err(e) => {
+                    warn!("Failed to enqueue OG image job for {crate_name}: {e}");
+                    errors += 1;
+                }
+            }
+        }
+
+        // Break if we've processed fewer than batch_size (last batch)
+        if batch_size < opts.batch_size {
+            break;
+        }
+
+        offset += opts.batch_size as i64;
+    }
+
+    info!("Jobs enqueued: {enqueued}");
+    info!("Errors: {errors}");
+
+    if errors > 0 {
+        warn!("Some jobs failed to enqueue. Check logs above for details.");
+    }
+
+    Ok(())
+}

src/bin/crates-admin/main.rs

@@ -1,6 +1,7 @@
 #[macro_use]
 extern crate tracing;
 
+mod backfill_og_images;
 mod default_versions;
 mod delete_crate;
 mod delete_version;
@@ -17,6 +18,7 @@ mod yank_version;
 #[derive(clap::Parser, Debug)]
 #[command(name = "crates-admin")]
 enum Command {
+    BackfillOgImages(backfill_og_images::Opts),
     DeleteCrate(delete_crate::Opts),
     DeleteVersion(delete_version::Opts),
     Populate(populate::Opts),
@@ -46,6 +48,7 @@ async fn main() -> anyhow::Result<()> {
     span.record("command", tracing::field::debug(&command));
 
     match command {
+        Command::BackfillOgImages(opts) => backfill_og_images::run(opts).await,
         Command::DeleteCrate(opts) => delete_crate::run(opts).await,
         Command::DeleteVersion(opts) => delete_version::run(opts).await,
         Command::Populate(opts) => populate::run(opts).await,