use std::fmt::Write as _;
use super::{
Edits,
fields::{FieldType, parse_fields},
mailer::pascal,
next_migration_path, register_routes,
schedule::{FREE_CRONS, add_task, parse_cron},
};
use crate::{
CliResult,
names::{RUST_KEYWORDS, is_identifier},
output::CliError,
project::Project,
secret,
};
const SWEEP: &str = "*/5 * * * *";
pub fn external_job(project: &Project, name: &str, specs: &[String], sweep: Option<&str>, realtime: bool) -> CliResult {
let name = name.strip_suffix("_jobs").or_else(|| name.strip_suffix("_job")).unwrap_or(name);
if !is_identifier(name) || name.chars().any(|c| c.is_ascii_uppercase()) || RUST_KEYWORDS.contains(&name) {
return Err(CliError::new(format!("invalid external job name `{name}`"))
.hint("use snake_case, e.g. `ocre g external_job upscale video_id:integer`"));
}
let fields = parse_fields(specs)?;
if let Some(field) = fields.iter().find(|field| {
matches!(field.ty, FieldType::Attachment | FieldType::Enum | FieldType::RichText) || field.polymorphic.is_some()
}) {
return Err(CliError::new(format!("external job field `{}` cannot be of that type", field.name)).hint(
"inputs are sent as JSON: use strings, numbers, booleans, json or references (store files in R2 and \
pass their key or a presigned URL)",
));
}
if let Some(field) = fields.iter().find(|field| COLUMNS.iter().any(|(column, _)| *column == field.name)) {
return Err(CliError::new(format!("external job field `{}` is a column of the jobs table", field.name))
.hint("rename the field"));
}
let cron = parse_cron(sweep.unwrap_or(SWEEP))?;
let module = format!("{name}_jobs");
let upper = name.to_uppercase();
let command = std::iter::once(format!("ocre g external_job {name}"))
.chain(specs.iter().cloned())
.chain(sweep.map(|sweep| format!("--sweep \"{sweep}\"")))
.chain(realtime.then(|| "--realtime".to_owned()))
.collect::<Vec<_>>()
.join(" ");
let mut edits = Edits::new(project);
if edits.exists(&format!("src/{name}_webhook.rs")) {
return Err(CliError::new(format!("src/{name}_webhook.rs already answers /webhooks/{name}"))
.hint("pick another name for the external job"));
}
edits.create(&format!("src/{module}.rs"), module_rs(name, &upper, &fields, &command, realtime))?;
if realtime {
super::realtime::add_channel_prefix(&mut edits, &module, &command)?;
}
let mut columns = vec![
"id INTEGER PRIMARY KEY".to_owned(),
"status TEXT NOT NULL DEFAULT 'queued' CHECK (status IN ('queued', 'submitted', 'running', 'done', 'failed'))"
.to_owned(),
"external_id TEXT".to_owned(),
];
columns.extend(fields.iter().flat_map(|field| field.sql_columns()));
columns.extend(COLUMNS.iter().skip(3).map(|(column, sql)| format!("{column} {sql}")));
let sql = format!(
"-- Generated by `{command}`: jobs run by an external service.\nCREATE TABLE {module} (\n {}\n);\n\nCREATE INDEX index_{module}_on_status_and_updated_at ON {module} (status, updated_at);\n",
columns.join(",\n ")
);
let path = next_migration_path(&edits, &format!("create_{module}"))?;
edits.create(&path, sql)?;
let count = add_task(&mut edits, &format!("{module}_sweep"), &cron, sweep_rs(&module, &cron, &command))?;
if let Some(vars) = edits.read(".dev.vars")? {
let vars = if vars.is_empty() || vars.ends_with('\n') { vars } else { format!("{vars}\n") };
let mut added = String::new();
for (var, value) in [
(format!("{upper}_URL"), "http://localhost:8787/fake".to_owned()),
(format!("{upper}_SECRET"), secret::generate()),
("APP_URL".to_owned(), "http://localhost:8787".to_owned()),
] {
if !vars.lines().any(|line| line.split('=').next().is_some_and(|key| key.trim() == var)) {
writeln!(added, "{var}={value}").expect("writing to a String");
}
}
edits.update(".dev.vars", format!("{vars}{added}"));
}
register_routes(&mut edits, &module)?;
let mut report = edits.apply("generate external_job")?;
report.next = vec![
"ocre migrate".to_owned(),
format!(
"set {upper}_URL in .dev.vars to the service's endpoint, and adapt `request_body` and `handle_event` in \
src/{module}.rs to its API"
),
format!("start a job from a handler: {module}::start(&ctx, ...).await?"),
format!(
"production: {upper}_URL and APP_URL in worker.env of cloudflare.config.ts; {upper}_SECRET (and \
{upper}_TOKEN for a bearer API key) in .prod.vars, then `ocre secrets push {upper}_SECRET --file .prod.vars`"
),
];
if count > FREE_CRONS {
report.next.push(format!(
"this app now has {count} crons; the free plan allows {FREE_CRONS} per account: run several tasks from one cron"
));
}
Ok(report)
}
const COLUMNS: [(&str, &str); 11] = [
("id", ""),
("status", ""),
("external_id", ""),
("public_id", "TEXT NOT NULL UNIQUE"),
("progress", "INTEGER NOT NULL DEFAULT 0"),
("result", "TEXT"),
("error", "TEXT"),
("attempts", "INTEGER NOT NULL DEFAULT 0"),
("token", "TEXT NOT NULL"),
("created_at", "INTEGER NOT NULL"),
("updated_at", "INTEGER NOT NULL"),
];
fn module_rs(name: &str, upper: &str, fields: &[super::fields::Field], command: &str, realtime: bool) -> String {
let module = format!("{name}_jobs");
let model = format!("{}Job", pascal(name));
let mut struct_fields = String::new();
let mut params = Vec::new();
let mut input = Vec::new();
for field in fields {
let attribute = match (field.ty, field.optional) {
(FieldType::Boolean, _) => " #[serde(deserialize_with = \"ocre::bool_from_sql\")]\n",
(FieldType::Json, false) => " #[serde(deserialize_with = \"ocre::json_from_sql\")]\n",
(FieldType::Json, true) => " #[serde(default, deserialize_with = \"ocre::optional_json_from_sql\")]\n",
_ => "",
};
writeln!(struct_fields, "{attribute} pub {}: {},", field.name, field.column_type())
.expect("writing to a String");
params.push(format!("{}: {}", field.name, field.column_type()));
input.push(format!("\"{0}\": job.{0}", field.name));
}
let names: Vec<&str> = fields.iter().map(|field| field.name.as_str()).collect();
let column_list = names.iter().map(|n| format!("{n}, ")).collect::<String>();
let placeholders: String = (0..names.len()).map(|i| format!("?{}, ", i + 1)).collect();
let n = names.len();
let insert = format!(
"INSERT INTO {module} ({column_list}public_id, token, created_at, updated_at) VALUES ({placeholders}?{}, ?{}, ?{}, ?{}) RETURNING *",
n + 1,
n + 2,
n + 3,
n + 3
);
let insert_params = names.iter().map(|n| format!("{n}, ")).collect::<String>()
+ "ocre::token::public_id(), ocre::token::generate(), ocre::now()";
let (live_imports, live_route, changed) = if realtime {
(
"\n response::Html,",
format!("\n .route(\"/{module}/{{public_id}}/progress\", get(progress))"),
format!(
r##"/// Runs after every change of a job (submitted, an event, failed by the
/// sweep): its progress bar moves in every page showing it. Add more: notify
/// the user, enqueue the next job.
async fn changed(ctx: &Ctx, job: &{model}) -> Result<()> {{
// Best effort: Ocre logs a failed broadcast; the job goes on.
ocre::realtime::broadcast(ctx, &channel(job), &progress_html(job)).await.ok();
Ok(())
}}
/// The realtime channel of a job, named by its public id.
fn channel(job: &{model}) -> String {{
format!("{module}:{{}}", job.public_id)
}}
/// The job's progress bar (`ocre::helpers::progress_bar`); a broadcast of
/// a new one replaces it, as they share the element id.
pub fn progress_html(job: &{model}) -> String {{
let label = match (job.status.as_str(), &job.error) {{
("queued", _) => "Waiting".to_owned(),
("submitted", _) => "Sent".to_owned(),
("running", _) => "Running".to_owned(),
("done", _) => "Done".to_owned(),
(_, Some(error)) => format!("Failed: {{error}}"),
_ => "Failed".to_owned(),
}};
ocre::helpers::progress_bar(&format!("{module}_{{}}", job.public_id), job.progress, &label)
}}
/// `GET /{module}/<public_id>/progress`: the job's progress bar, kept up to
/// date over a WebSocket. Load it in a page that has htmx and its `ws`
/// extension: `<div hx-get="/{module}/<public_id>/progress" hx-trigger="load"></div>`.
async fn progress(State(ctx): State<Ctx>, Path(public_id): Path<String>) -> Result<Html<String>> {{
let job = find_by_public_id(&ctx, &public_id).await?.ok_or(Error::NotFound)?;
Ok(Html(format!(
r#"<div hx-ext="ws" ws-connect="/realtime/{{}}">{{}}</div>"#,
channel(&job),
progress_html(&job)
)))
}}
"##,
module = module
),
)
} else {
(
"",
String::new(),
format!(
r#"/// Runs after every change of a job (submitted, an event, failed by the
/// sweep), e.g. `ocre::realtime::broadcast(ctx, &format!("{name}:{{}}", job.public_id), &html)`
/// (or generate with `--realtime`).
async fn changed(_ctx: &Ctx, _job: &{model}) -> Result<()> {{
Ok(())
}}
"#
),
)
};
format!(
r##"//! `{name}` jobs, run by an external service (a GPU on RunPod or Modal, a
//! Cloudflare Container, any HTTP API). Generated by `{command}`.
//!
//! A job is a row of `{module}`:
//!
//! ```text
//! queued ──submit──▶ submitted ──event──▶ running ──event──▶ done
//! │ ▲ │ │
//! │ └─ sweep (retry) └──────────────────┴──event or sweep (no news)──▶ failed
//! └── MAX_ATTEMPTS submissions failed ──▶ failed
//! ```
//!
//! - `start(&ctx, ...)` inserts a `queued` job and submits it: a POST of
//! `request_body(...)` to `{upper}_URL`, signed (`X-Signature`, HMAC-SHA256
//! of the body with `{upper}_SECRET`), with `Authorization: Bearer
//! {upper}_TOKEN` when that secret is set. The body tells the service where
//! to report: `<APP_URL>/webhooks/{name}/<id>?token=<the job's token>`.
//! - The service POSTs events there: JSON with a `status` (`running`,
//! `done`, `failed`, or RunPod's `IN_PROGRESS`, `COMPLETED`, `FAILED`...),
//! and optionally `progress` (0-100), `output` or `result`, `error`. An
//! event is accepted when signed (`X-Signature`) or when it carries the
//! job's token (services such as RunPod cannot sign). Events only move a
//! job forward: a repeated or late event changes nothing.
//! - `sweep` (src/schedules/{module}_sweep.rs) resubmits jobs whose
//! submission failed, and fails jobs without news for `STALE_AFTER` seconds.
//! - `changed(&ctx, &job)` runs after every change: notify the user, enqueue
//! the next job, broadcast the progress (see the realtime guide).
// `start`, `find_by_public_id`... are generated before a handler calls them.
#![allow(dead_code)]
use axum::{{
Json, Router,
body::Bytes,
extract::{{Path, Query, State}},
http::HeaderMap,{live_imports}
routing::{routing},
}};
use ocre::{{
Ctx, Error, Result, params,
serde_json::{{self, Value, json}},
webhooks,
}};
use serde::{{Deserialize, Serialize}};
/// Submissions tried before a job fails.
pub const MAX_ATTEMPTS: i64 = 3;
/// Seconds without news (an event or a submission) before `sweep` fails a job.
pub const STALE_AFTER: i64 = 3600;
/// Seconds a failed submission waits before `sweep` tries again.
pub const RETRY_AFTER: i64 = 60;
/// A job, as stored in `{module}`. Times are Unix seconds.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct {model} {{
pub id: i64,
/// `queued`, `submitted`, `running`, `done` or `failed`.
pub status: String,
/// The service's id for the job, from its answer to the submission.
pub external_id: Option<String>,
/// Random, for pages and channels: `find_by_public_id`.
pub public_id: String,
{struct_fields} /// 0 to 100, as the service reports it.
pub progress: i64,
/// The service's output, as JSON text.
pub result: Option<String>,
pub error: Option<String>,
/// Submissions tried.
pub attempts: i64,
/// Proves an event comes from the service this job was sent to.
#[serde(skip_serializing)]
pub token: String,
pub created_at: i64,
pub updated_at: i64,
}}
impl {model} {{
/// Done or failed: no event changes it any more.
pub fn is_finished(&self) -> bool {{
self.status == "done" || self.status == "failed"
}}
}}
pub fn routes() -> Router<Ctx> {{
Router::new().route("/webhooks/{name}/{{id}}", post(receive)){live_route}
}}
/// The job `id`, or `Error::NotFound`.
pub async fn find(ctx: &Ctx, id: i64) -> Result<{model}> {{
ctx.db()?.first("SELECT * FROM {module} WHERE id = ?1", params![id]).await?.ok_or(Error::NotFound)
}}
/// The job with this public id, if any.
pub async fn find_by_public_id(ctx: &Ctx, public_id: &str) -> Result<Option<{model}>> {{
ctx.db()?.first("SELECT * FROM {module} WHERE public_id = ?1", params![public_id]).await
}}
/// Creates a job and submits it. The job comes back `submitted`, or
/// `queued` with the error when the service could not take it (the sweep
/// tries again).
pub async fn start(ctx: &Ctx{params}) -> Result<{model}> {{
let job: {model} = ctx
.db()?
.first("{insert}", params![{insert_params}])
.await?
.ok_or_else(|| Error::internal("INSERT ... RETURNING returned no row"))?;
submit(ctx, job).await
}}
/// What the service receives: the job's input and where to report. RunPod
/// reads `input` and `webhook`; adapt it to the service's API.
fn request_body(job: &{model}, webhook: &str) -> Value {{
json!({{ "id": job.id, "input": {{ {input} }}, "webhook": webhook }})
}}
/// Sends a `queued` job to the service; returns the job as it is afterwards.
pub async fn submit(ctx: &Ctx, job: {model}) -> Result<{model}> {{
let env = ctx.env();
let url = env.var("{upper}_URL").map_err(|_| Error::internal("{upper}_URL is not set"))?.to_string();
let secret = ctx.secret("{upper}_SECRET").await?;
let bearer = ctx.secret("{upper}_TOKEN").await.ok();
let webhook = format!("{{}}?token={{}}", ocre::mail::url(ctx, &format!("/webhooks/{name}/{{}}", job.id))?, job.token);
let answer = webhooks::post_signed(&url, secret.as_bytes(), bearer.as_deref(), &request_body(&job, &webhook)).await;
let db = ctx.db()?;
let now = ocre::now();
let updated: Option<{model}> = match answer {{
Ok(answer) if (200..300).contains(&answer.status) => {{
let external_id = serde_json::from_str::<Value>(&answer.text)
.ok()
.and_then(|value| value["id"].as_str().map(str::to_owned));
let sql = "UPDATE {module} SET status = 'submitted', external_id = ?2, attempts = attempts + 1, \
error = NULL, updated_at = ?3 WHERE id = ?1 AND status = 'queued' RETURNING *";
db.first(sql, params![job.id, external_id, now]).await?
}}
failed => {{
let error = match failed {{
Ok(answer) => format!("{upper}_URL answered {{}}: {{}}", answer.status, answer.text.chars().take(200).collect::<String>()),
Err(err) => err.to_string(),
}};
let sql = "UPDATE {module} SET status = CASE WHEN attempts + 1 >= ?2 THEN 'failed' ELSE 'queued' END, \
attempts = attempts + 1, error = ?3, updated_at = ?4 WHERE id = ?1 AND status = 'queued' RETURNING *";
db.first(sql, params![job.id, MAX_ATTEMPTS, error, now]).await?
}}
}};
match updated {{
Some(job) => {{
changed(ctx, &job).await?;
Ok(job)
}}
// An event moved it on meanwhile.
None => find(ctx, job.id).await,
}}
}}
#[derive(Deserialize)]
struct Auth {{
token: Option<String>,
}}
async fn receive(
State(ctx): State<Ctx>,
Path(id): Path<i64>,
Query(auth): Query<Auth>,
headers: HeaderMap,
body: Bytes,
) -> Result<Json<Value>> {{
let job = find(&ctx, id).await?;
let secret = ctx.secret("{upper}_SECRET").await?;
let signature = headers.get("x-signature").and_then(|value| value.to_str().ok());
let signed = signature.is_some_and(|signature| webhooks::verify(secret.as_bytes(), &body, signature).is_ok());
let tokened =
auth.token.is_some_and(|token| ocre::token::constant_time_eq(token.as_bytes(), job.token.as_bytes()));
if !signed && !tokened {{
return Err(Error::Unauthorized);
}}
let event: Value = serde_json::from_slice(&body).map_err(|_| Error::bad_request("the body is not JSON"))?;
handle_event(&ctx, job, &event).await?;
Ok(Json(json!({{ "ok": true }})))
}}
/// Applies an event of the service to the job. Only forward moves count, so
/// a repeated or late event changes nothing.
async fn handle_event(ctx: &Ctx, job: {model}, event: &Value) -> Result<()> {{
let reported = event["status"].as_str().unwrap_or_default();
let status = match reported.to_ascii_lowercase().as_str() {{
"running" | "started" | "processing" | "in_progress" | "in_queue" => "running",
"done" | "completed" | "succeeded" | "success" => "done",
"failed" | "error" | "cancelled" | "canceled" | "timed_out" => "failed",
_ => return Err(Error::bad_request(format!("unknown status `{{reported}}`"))),
}};
let progress = if status == "done" {{ Some(100) }} else {{ event["progress"].as_i64().map(|p| p.clamp(0, 100)) }};
let result = event.get("output").or_else(|| event.get("result")).filter(|value| !value.is_null()).map(Value::to_string);
let error = match (status, event["error"].as_str()) {{
(_, Some(error)) => Some(error.to_owned()),
("failed", None) => Some(format!("the service reported `{{reported}}`")),
_ => None,
}};
let sql = "UPDATE {module} SET status = ?2, progress = MAX(progress, COALESCE(?3, progress)), \
result = COALESCE(?4, result), error = COALESCE(?5, error), updated_at = ?6 \
WHERE id = ?1 AND status IN ('queued', 'submitted', 'running') RETURNING *";
let updated: Option<{model}> =
ctx.db()?.first(sql, params![job.id, status, progress, result, error, ocre::now()]).await?;
if let Some(job) = updated {{
changed(ctx, &job).await?;
}}
Ok(())
}}
/// Run by the Cron Trigger: fails jobs without news for `STALE_AFTER`
/// seconds and resubmits queued ones (10 per run: each is a subrequest).
pub async fn sweep(ctx: &Ctx) -> Result<()> {{
let db = ctx.db()?;
let now = ocre::now();
let sql = "UPDATE {module} SET status = 'failed', error = 'no news from the service for ' || ?2 || ' s', \
updated_at = ?1 WHERE status IN ('submitted', 'running') AND updated_at < ?1 - ?2 RETURNING *";
let stale: Vec<{model}> = db.all(sql, params![now, STALE_AFTER]).await?;
for job in &stale {{
changed(ctx, job).await?;
}}
let sql = "SELECT * FROM {module} WHERE status = 'queued' AND updated_at < ?1 ORDER BY id LIMIT 10";
let queued: Vec<{model}> = db.all(sql, params![now - RETRY_AFTER]).await?;
for job in queued {{
submit(ctx, job).await?;
}}
Ok(())
}}
{changed}
"##,
input = input.join(", "),
routing = if realtime { "{get, post}" } else { "post" },
params = params.iter().map(|param| format!(", {param}")).collect::<String>(),
)
}
fn sweep_rs(module: &str, cron: &str, command: &str) -> String {
format!(
r#"//! Sweeps the jobs of src/{module}.rs at `{cron}` (UTC). Generated by `{command}`.
//! Run it now while `ocre dev` runs: `ocre schedules run {module}_sweep`.
use ocre::{{Ctx, Result}};
pub async fn run(ctx: &Ctx) -> Result<()> {{
crate::{module}::sweep(ctx).await
}}
"#
)
}