feature: connect batch etl to trieve

batch-etl/Cargo.toml

@@ -2,6 +2,19 @@
 name = "batch-etl"
 version = "0.1.0"
 edition = "2021"
+default-run = "batch-etl-server"
+
+[[bin]]
+name = "batch-etl-server"
+path = "src/main.rs"
+
+[[bin]]
+name = "create-job"
+path = "src/workers/create-job-worker.rs"
+
+[[bin]]
+name = "check-completion"
+path = "src/cron-jobs/check-completion.rs"
 
 [dependencies]
 utoipa-redoc = { version = "5.0.0", features = ["actix-web"] }
@@ -26,6 +39,7 @@ time = "0.3.37"
 rust-s3 = "0.35.1"
 reqwest = "0.12.9"
 jsonl = "4.0.1"
+broccoli_queue = "0.1.2"
 
 [features]
 default = []

batch-etl/ch_migrations/1734595053_initial_tables/down.sql

@@ -1,3 +1,4 @@
 DROP TABLE IF EXISTS inputs;
 DROP TABLE IF EXISTS jobs;
-DROP TABLE IF EXISTS schemas;
+DROP TABLE IF EXISTS schemas;
+DROP TABLE IF EXISTS batches;

batch-etl/ch_migrations/1734595053_initial_tables/up.sql

@@ -13,12 +13,10 @@ CREATE TABLE IF NOT EXISTS jobs (
     id String,
     schema_id String,
     input_id String,
-    status String,
-    batch_id String,
-    output_id String,
+    webhook_url String,
     created_at DateTime,
     updated_at DateTime
-) ENGINE = ReplacingMergeTree(updated_at)
+) ENGINE = MergeTree()
 ORDER BY (schema_id, id)
 PARTITION BY
     (schema_id);
@@ -29,3 +27,16 @@ CREATE TABLE IF NOT EXISTS inputs (
     updated_at DateTime
 ) ENGINE = MergeTree()
 ORDER BY (id);
+
+CREATE TABLE IF NOT EXISTS batches (
+    batch_id String,
+    job_id String,
+    output_id String,
+    status String,
+    created_at DateTime,
+    updated_at DateTime
+) ENGINE = ReplacingMergeTree(updated_at)
+ORDER BY (job_id, batch_id)
+PARTITION BY
+    (job_id);
+

batch-etl/ch_migrations/chm.toml

@@ -1,4 +1,4 @@
-url = "http://localhost:8123"
+url = "http://localhost:8124"
 user = "clickhouse"
 password = "password"
 database = "default"

batch-etl/docker-compose.yml

@@ -8,7 +8,7 @@ services:
       timeout: 5s
       retries: 10
     ports:
-      - "6379:6379"
+      - "6380:6379"
     volumes:
       - redis-data:/data
     networks:
@@ -18,8 +18,8 @@ services:
   s3:
     image: minio/minio:RELEASE.2023-09-27T15-22-50Z
     ports:
-      - 9000:9000
-      - 42625:42625
+      - 9004:9000
+      - 42626:42625
     environment:
       - MINIO_ROOT_USER=${MINIO_ROOT_USER}
       - MINIO_ROOT_PASSWORD=${MINIO_ROOT_PASSWORD}
@@ -64,9 +64,9 @@ services:
     volumes:
       - clickhouse-data:/var/lib/clickhouse
     ports:
-      - "8123:8123"
-      - "9001:9000"
-      - "9009:9009"
+      - "8124:8123"
+      - "9002:9000"
+      - "9010:9009"
     networks:
       - app-network
 

batch-etl/src/cron-jobs/check-completion.rs

@@ -0,0 +1,38 @@
+use batch_etl::operators::{
+    batch::{get_batch_output, get_pending_batches},
+    job::send_webhook,
+};
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    dotenvy::dotenv().ok();
+    env_logger::builder()
+        .target(env_logger::Target::Stdout)
+        .filter_level(log::LevelFilter::Info)
+        .init();
+
+    let clickhouse_client = clickhouse::Client::default()
+        .with_url(std::env::var("CLICKHOUSE_URL").unwrap_or("http://localhost:8123".to_string()))
+        .with_user(std::env::var("CLICKHOUSE_USER").unwrap_or("default".to_string()))
+        .with_password(std::env::var("CLICKHOUSE_PASSWORD").unwrap_or("".to_string()))
+        .with_database(std::env::var("CLICKHOUSE_DATABASE").unwrap_or("default".to_string()))
+        .with_option("async_insert", "1")
+        .with_option("wait_for_async_insert", "0");
+
+    let pending_batches = get_pending_batches(&clickhouse_client).await.unwrap();
+    log::info!("Pending batches: {:?}", pending_batches);
+
+    for (batch, id, webhook_url) in pending_batches {
+        log::info!("Processing batch: {:?}", batch);
+        let batch_url = get_batch_output(&clickhouse_client, batch).await.unwrap();
+        log::info!("Batch URL: {:?}", batch_url);
+        if let Some(webhook_url) = webhook_url {
+            if let Some(batch_url) = batch_url {
+                log::info!("Sending webhook to: {:?}", webhook_url);
+                send_webhook(webhook_url, id, batch_url).await.unwrap();
+            }
+        }
+    }
+
+    Ok(())
+}

batch-etl/src/lib.rs

@@ -1,7 +1,9 @@
 use actix_web::{
+    middleware::Logger,
     web::{self, PayloadConfig},
     App, HttpServer,
 };
+use broccoli_queue::queue::BroccoliQueue;
 use chm::tools::migrations::{run_pending_migrations, SetupArgs};
 use errors::custom_json_error_handler;
 use routes::{
@@ -131,6 +133,17 @@ pub async fn main() -> std::io::Result<()> {
         .limit(134200000)
         .error_handler(custom_json_error_handler);
 
+    let broccoli_queue = BroccoliQueue::builder(redis_url)
+        .pool_connections(redis_connections.try_into().unwrap())
+        .build()
+        .await
+        .map_err(|e| {
+            std::io::Error::new(
+                std::io::ErrorKind::Other,
+                format!("Failed to create broccoli queue {:?}", e),
+            )
+        })?;
+
     HttpServer::new(move || {
         App::new()
             .into_utoipa_app()
@@ -139,6 +152,7 @@ pub async fn main() -> std::io::Result<()> {
             .app_data(json_cfg.clone())
             .app_data(web::Data::new(redis_pool.clone()))
             .app_data(web::Data::new(clickhouse_client.clone()))
+            .app_data(web::Data::new(broccoli_queue.clone()))
             .service(utoipa_actix_web::scope("/api/schema").configure(|config| {
                 config.service(create_schema).service(get_schema);
             }))
@@ -154,6 +168,13 @@ pub async fn main() -> std::io::Result<()> {
             }))
             .openapi_service(|api| Redoc::with_url("/redoc", api))
             .into_app()
+            .wrap(
+                // Set up logger, but avoid logging hot status endpoints
+                Logger::new("%r %s %b %{Referer}i %{User-Agent}i %T %{TR-Dataset}i")
+                    .exclude("/")
+                    .exclude("/api/health")
+                    .exclude("/metrics"),
+            )
     })
     .bind(("0.0.0.0", 8082))?
     .run()