jobs.rs

use crate::queries::api_queries::{DbCheckpoint, DbLogMessage, DbPipelineJob};
use anyhow::Context;
use arroyo_rpc::api_types::checkpoints::{
    Checkpoint, CheckpointType, JobCheckpointSpan, OperatorCheckpointGroup, SubtaskCheckpointGroup,
};
use arroyo_rpc::api_types::pipelines::{JobLogLevel, JobLogMessage, OutputData, StopType};
use arroyo_rpc::api_types::{
    CheckpointCollection, JobCollection, JobLogMessageCollection,
    OperatorCheckpointGroupCollection, PaginationQueryParams,
};
use arroyo_rpc::grpc::api::OperatorCheckpointDetail;
use arroyo_rpc::public_ids::{IdTypes, generate_id};
use arroyo_rpc::{LeaderContext, StateContext, get_event_spans, grpc, job_status_client};
use axum::Json;
use axum::extract::{Path, Query, State};
use axum::response::sse::{Event, Sse};
use futures_util::stream::Stream;
use std::convert::Infallible;
use std::str::FromStr;
use std::{collections::HashMap, time::Duration};
use tokio_stream::StreamExt as _;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Code, Request};
use tracing::info;

const PREVIEW_TTL: Duration = Duration::from_secs(60);

use crate::pipelines::{query_job_by_pub_id, query_pipeline_by_pub_id};
use crate::rest::AppState;
use crate::rest_utils::{
    BearerAuth, ErrorResp, PipelineJobCheckpointPath, PipelineJobPath, authenticate, bad_request,
    conflict, internal_server_error, log_and_map, not_found, paginate_results,
    validate_pagination_params,
};
use crate::types::public::LogLevel;
use crate::{AuthData, queries::api_queries, to_micros, types::public};
use arroyo_rpc::config::config;
use arroyo_rpc::controller_client;
use arroyo_rpc::errors::ErrorDomain;
use arroyo_rpc::grpc::rpc::job_status_grpc_client::JobStatusGrpcClient;
use arroyo_rpc::grpc::rpc::{GetCheckpointDetailsReq, GetJobCheckpointsReq};
use arroyo_rpc::identity::InjectWorkerId;
use cornucopia_async::DatabaseSource;
use http::StatusCode;
use tonic::codegen::InterceptedService;
use tonic::transport::Channel;

type LeaderCheckpointClient = JobStatusGrpcClient<InterceptedService<Channel, InjectWorkerId>>;

async fn fetch_from_leader<
    R,
    F: Future<Output = Result<R, tonic::Status>>,
    T: FnOnce(LeaderCheckpointClient, u64) -> F,
>(
    leader_context: LeaderContext,
    req: T,
) -> Result<R, ErrorResp> {
    let client = job_status_client(
        "api",
        &config().api.tls,
        leader_context.worker_id,
        leader_context.rpc_address,
        None,
    )
    .await
    .map_err(|_| ErrorResp {
        status_code: StatusCode::BAD_GATEWAY,
        message: "the leader is not currently available; try again later".to_string(),
    })?;

    req(client, leader_context.generation)
        .await
        .map_err(|e| match e.code() {
            Code::FailedPrecondition | Code::NotFound => bad_request(e.message().to_string()),
            Code::Unavailable => ErrorResp {
                status_code: StatusCode::BAD_GATEWAY,
                message: "the leader is not currently available; try again later".to_string(),
            },
            _ => log_and_map(e),
        })
}

fn operator_checkpoint_groups(
    operator_details: HashMap<String, OperatorCheckpointDetail>,
) -> Vec<OperatorCheckpointGroup> {
    let mut operators = vec![];

    operator_details
        .iter()
        .for_each(|(operator_id, operator_details)| {
            let mut operator_bytes = 0;
            let mut subtasks = vec![];

            operator_details
                .tasks
                .iter()
                .for_each(|(subtask_index, subtask_details)| {
                    operator_bytes += subtask_details.bytes.unwrap_or(0);
                    subtasks.push(SubtaskCheckpointGroup {
                        index: *subtask_index,
                        bytes: subtask_details.bytes.unwrap_or(0),
                        event_spans: get_event_spans(subtask_details).into(),
                    });
                });

            operators.push(OperatorCheckpointGroup {
                operator_id: operator_id.to_string(),
                bytes: operator_bytes,
                started_metadata_write: operator_details.started_metadata_write,
                finish_time: operator_details.finish_time,
                subtasks,
            });
        });

    operators
}

#[allow(clippy::too_many_arguments)]
pub(crate) async fn create_job(
    pipeline_name: &str,
    pipeline_id: i64,
    checkpoint_interval: Duration,
    preview: bool,
    auth: &AuthData,
    db: &DatabaseSource,
    env_vars: HashMap<String, String>,
    scheduler_config: serde_json::Value,
) -> Result<String, ErrorResp> {
    let checkpoint_interval = if preview {
        Duration::from_secs(24 * 60 * 60)
    } else {
        checkpoint_interval
    };

    if checkpoint_interval < Duration::from_secs(1)
        || checkpoint_interval > Duration::from_secs(24 * 60 * 60)
    {
        return Err(bad_request(
            "Checkpoint_interval_micros must be between 1 second and 1 day.".to_string(),
        ));
    }

    let running_jobs = api_queries::fetch_get_jobs(&db.client().await?, &auth.organization_id)
        .await?
        .iter()
        .filter(|j| {
            j.stop == public::StopMode::none
                && !j
                    .state
                    .as_ref()
                    .map(|s| s == "Failed" || s == "Finished")
                    .unwrap_or(false)
        })
        .count();

    if running_jobs > auth.org_metadata.max_running_jobs as usize {
        let message = format!("You have exceeded the maximum number
            of running jobs in your plan ({}). Stop an existing job or contact support@arroyo.systems for
            an increase", auth.org_metadata.max_running_jobs);

        return Err(bad_request(message));
    }

    let job_id = generate_id(IdTypes::JobConfig);

    let env_vars_json =
        serde_json::to_value(&env_vars).expect("HashMap<String, String> is always serializable");

    // TODO: handle chance of collision in ids
    api_queries::execute_create_job(
        &db.client().await?,
        &job_id,
        &auth.organization_id,
        &pipeline_name,
        &auth.user_id,
        &pipeline_id,
        &(checkpoint_interval.as_micros() as i64),
        &(if preview {
            Some(PREVIEW_TTL.as_micros() as i64)
        } else {
            None
        }),
        &env_vars_json,
        &scheduler_config,
    )
    .await?;

    api_queries::execute_create_job_status(
        &db.client().await?,
        &generate_id(IdTypes::JobStatus),
        &job_id,
        &auth.organization_id,
    )
    .await?;

    Ok(job_id)
}

fn replaceable_job(jobs: Vec<(String, String)>) -> Result<String, ErrorResp> {
    let count = jobs.len();
    let Some((job_id, state)) = jobs.into_iter().next() else {
        return Err(not_found("Job for pipeline"));
    };

    if count != 1 {
        return Err(internal_server_error(format!(
            "expected one job for pipeline, found {count}"
        )));
    }

    if state != "Stopped" && state != "Failed" {
        return Err(conflict(format!(
            "cannot restart job {job_id} without state while it is in state {state}; stop the job first"
        )));
    }

    Ok(job_id)
}

/// Replaces a pipeline's single terminal job with a fresh job.
///
/// This is used for restart-without-state in leader mode. Keeping the replacement in a single
/// transaction preserves the current one-job-per-pipeline assumption for all other API queries.
pub(crate) async fn replace_job_without_state(
    db: &DatabaseSource,
    pipeline_pub_id: &str,
    auth: &AuthData,
) -> Result<String, ErrorResp> {
    let new_job_id = generate_id(IdTypes::JobConfig);
    let new_status_id = generate_id(IdTypes::JobStatus);

    match db {
        DatabaseSource::Postgres(pool) => {
            let mut client = pool.get().await.map_err(log_and_map)?;
            let transaction = client.transaction().await.map_err(log_and_map)?;

            // Lock the pipeline first. This serializes concurrent replacement requests before
            // either request resolves the pipeline's current singleton job.
            let pipeline = transaction
                .query_opt(
                    "SELECT id FROM pipelines \
                     WHERE pub_id = $1 AND organization_id = $2 \
                     FOR UPDATE",
                    &[&pipeline_pub_id, &auth.organization_id],
                )
                .await
                .map_err(log_and_map)?
                .ok_or_else(|| not_found("Pipeline"))?;
            let pipeline_id: i64 = pipeline.get(0);

            let jobs = transaction
                .query(
                    "SELECT c.id, COALESCE(s.state, 'Created') \
                     FROM job_configs c \
                     INNER JOIN job_statuses s ON c.id = s.id \
                     WHERE c.pipeline_id = $1 AND c.organization_id = $2 \
                     FOR UPDATE OF c, s",
                    &[&pipeline_id, &auth.organization_id],
                )
                .await
                .map_err(log_and_map)?
                .into_iter()
                .map(|row| (row.get(0), row.get(1)))
                .collect();
            let old_job_id = replaceable_job(jobs)?;

            let inserted = transaction
                .execute(
                    "INSERT INTO job_configs \
                     (id, organization_id, pipeline_name, created_by, updated_by, updated_at, \
                      ttl_micros, stop, pipeline_id, parallelism_overrides, \
                      checkpoint_interval_micros, env_vars, scheduler_config) \
                     SELECT $1, organization_id, pipeline_name, created_by, $2, CURRENT_TIMESTAMP, \
                            ttl_micros, 'none', pipeline_id, parallelism_overrides, \
                            checkpoint_interval_micros, env_vars, scheduler_config \
                     FROM job_configs WHERE id = $3",
                    &[&new_job_id, &auth.user_id, &old_job_id],
                )
                .await
                .map_err(log_and_map)?;
            if inserted != 1 {
                return Err(internal_server_error("failed to clone job configuration"));
            }

            transaction
                .execute(
                    "INSERT INTO job_statuses (pub_id, id, organization_id) VALUES ($1, $2, $3)",
                    &[&new_status_id, &new_job_id, &auth.organization_id],
                )
                .await
                .map_err(log_and_map)?;

            // Delete children explicitly because the SQLite checkpoints table does not have the
            // same foreign-key cascade as PostgreSQL.
            transaction
                .execute("DELETE FROM checkpoints WHERE job_id = $1", &[&old_job_id])
                .await
                .map_err(log_and_map)?;
            transaction
                .execute(
                    "DELETE FROM job_log_messages WHERE job_id = $1",
                    &[&old_job_id],
                )
                .await
                .map_err(log_and_map)?;
            transaction
                .execute("DELETE FROM job_statuses WHERE id = $1", &[&old_job_id])
                .await
                .map_err(log_and_map)?;
            transaction
                .execute("DELETE FROM job_configs WHERE id = $1", &[&old_job_id])
                .await
                .map_err(log_and_map)?;

            transaction.commit().await.map_err(log_and_map)?;
        }
        DatabaseSource::Sqlite(connection) => {
            let mut connection = connection.lock().map_err(log_and_map)?;
            let transaction = connection.transaction().map_err(log_and_map)?;

            // Beginning a SQLite transaction acquires the connection mutex for this entire block,
            // so resolving and replacing the singleton job cannot interleave with another request.
            let pipeline_id = transaction
                .query_row(
                    "SELECT id FROM pipelines WHERE pub_id = ?1 AND organization_id = ?2",
                    rusqlite::params![pipeline_pub_id, auth.organization_id],
                    |row| row.get::<_, i64>(0),
                )
                .map_err(|error| match error {
                    rusqlite::Error::QueryReturnedNoRows => not_found("Pipeline"),
                    error => log_and_map(error),
                })?;

            let jobs = {
                let mut statement = transaction
                    .prepare(
                        "SELECT c.id, COALESCE(s.state, 'Created') \
                         FROM job_configs c \
                         INNER JOIN job_statuses s ON c.id = s.id \
                         WHERE c.pipeline_id = ?1 AND c.organization_id = ?2",
                    )
                    .map_err(log_and_map)?;
                statement
                    .query_map(
                        rusqlite::params![pipeline_id, auth.organization_id],
                        |row| Ok((row.get(0)?, row.get(1)?)),
                    )
                    .map_err(log_and_map)?
                    .collect::<Result<Vec<_>, _>>()
                    .map_err(log_and_map)?
            };
            let old_job_id = replaceable_job(jobs)?;

            let inserted = transaction
                .execute(
                    "INSERT INTO job_configs \
                     (id, organization_id, pipeline_name, created_by, updated_by, updated_at, \
                      ttl_micros, stop, pipeline_id, parallelism_overrides, \
                      checkpoint_interval_micros, env_vars, scheduler_config) \
                     SELECT ?1, organization_id, pipeline_name, created_by, ?2, CURRENT_TIMESTAMP, \
                            ttl_micros, 'none', pipeline_id, parallelism_overrides, \
                            checkpoint_interval_micros, env_vars, scheduler_config \
                     FROM job_configs WHERE id = ?3",
                    rusqlite::params![new_job_id, auth.user_id, old_job_id],
                )
                .map_err(log_and_map)?;
            if inserted != 1 {
                return Err(internal_server_error("failed to clone job configuration"));
            }

            transaction
                .execute(
                    "INSERT INTO job_statuses (pub_id, id, organization_id) VALUES (?1, ?2, ?3)",
                    rusqlite::params![new_status_id, new_job_id, auth.organization_id],
                )
                .map_err(log_and_map)?;
            transaction
                .execute(
                    "DELETE FROM checkpoints WHERE job_id = ?1",
                    rusqlite::params![old_job_id],
                )
                .map_err(log_and_map)?;
            transaction
                .execute(
                    "DELETE FROM job_log_messages WHERE job_id = ?1",
                    rusqlite::params![old_job_id],
                )
                .map_err(log_and_map)?;
            transaction
                .execute(
                    "DELETE FROM job_statuses WHERE id = ?1",
                    rusqlite::params![old_job_id],
                )
                .map_err(log_and_map)?;
            transaction
                .execute(
                    "DELETE FROM job_configs WHERE id = ?1",
                    rusqlite::params![old_job_id],
                )
                .map_err(log_and_map)?;

            transaction.commit().map_err(log_and_map)?;
        }
    }

    Ok(new_job_id)
}

pub(crate) fn get_action(state: &str, running_desired: &bool) -> (String, Option<StopType>, bool) {
    enum Progress {
        InProgress,
        Stable,
    }

    use Progress::*;
    use StopType::*;

    let (a, s, p) = match (state, running_desired) {
        ("Created", true) => ("Stop", Some(Checkpoint), InProgress),
        ("Created", false) => ("Start", Some(None), Stable),

        ("Compiling", true) => ("Stop", Some(Checkpoint), InProgress),
        ("Compiling", false) => ("Stopping", Option::None, InProgress),

        ("Scheduling", true) => ("Stop", Some(Checkpoint), InProgress),
        ("Scheduling", false) => ("Stopping", Option::None, InProgress),

        ("Running", true) => ("Stop", Some(Checkpoint), Stable),
        ("Running", false) => ("Stopping", Option::None, InProgress),

        ("Rescaling", true) => ("Stop", Some(Checkpoint), InProgress),
        ("Rescaling", false) => ("Stopping", Option::None, InProgress),

        ("CheckpointStopping", true) => ("Force Stop", Some(Immediate), InProgress),
        ("CheckpointStopping", false) => ("Force Stop", Some(Immediate), InProgress),

        ("Recovering", true) => ("Stop", Some(Checkpoint), InProgress),
        ("Recovering", false) => ("Stopping", Option::None, InProgress),

        ("Restarting", true) => ("Stop", Some(Checkpoint), InProgress),
        ("Restarting", false) => ("Stopping", Option::None, InProgress),

        ("Stopping", true) => ("Stopping", Some(Checkpoint), InProgress),
        ("Stopping", false) => ("Stopping", Option::None, InProgress),

        ("Stopped", true) => ("Starting", Option::None, InProgress),
        ("Stopped", false) => ("Start", Some(None), Stable),

        ("Finishing", true) => ("Finishing", Option::None, InProgress),
        ("Finishing", false) => ("Finishing", Option::None, InProgress),

        ("Finished", true) => ("Finished", Option::None, Stable),
        ("Finished", false) => ("Finished", Option::None, Stable),

        ("Failed", true) => ("Failed", Option::None, Stable),
        ("Failed", false) => ("Start", Some(None), Stable),

        // Failing is a transient state during graceful shutdown before Failed
        ("Failing", true) => ("Failing", Option::None, InProgress),
        ("Failing", false) => ("Failing", Option::None, InProgress),

        _ => panic!("unhandled state {state}"),
    };

    let in_progress = match p {
        InProgress => true,
        Stable => false,
    };

    (a.to_string(), s, in_progress)
}

/// List a job's error messages
#[utoipa::path(
    get,
    path = "/v1/pipelines/{pipeline_id}/jobs/{job_id}/errors",
    tag = "jobs",
    params(
        ("pipeline_id" = String, Path, description = "Pipeline id"),
        ("job_id" = String, Path, description = "Job id"),
        ("starting_after" = Option<String>, Query, description = "Starting after"),
        ("limit" = Option<u32>, Query, description = "Limit"),
    ),
    responses(
        (status = 200, description = "Got job's error messages", body = JobLogMessageCollection),
    ),
)]
pub async fn get_job_errors(
    State(state): State<AppState>,
    bearer_auth: BearerAuth,
    Path(PipelineJobPath {
        id: pipeline_pub_id,
        job_id: job_pub_id,
    }): Path<PipelineJobPath>,
    query_params: Query<PaginationQueryParams>,
) -> Result<Json<JobLogMessageCollection>, ErrorResp> {
    let auth_data = authenticate(&state.database, bearer_auth).await?;
    let db = state.database.client().await?;

    let (starting_after, limit) =
        validate_pagination_params(query_params.starting_after.clone(), query_params.limit)?;

    query_job_by_pub_id(&pipeline_pub_id, &job_pub_id, &db, &auth_data).await?;

    let errors = api_queries::fetch_get_operator_errors(
        &db,
        &auth_data.organization_id,
        &job_pub_id,
        &starting_after.unwrap_or_default(),
        &(limit as i32),
    )
    .await
    .map_err(log_and_map)?
    .into_iter()
    .map(|m| m.into())
    .collect();

    let (errors, has_more) = paginate_results(errors, limit);

    Ok(Json(JobLogMessageCollection {
        data: errors,
        has_more,
    }))
}

impl From<DbLogMessage> for JobLogMessage {
    fn from(val: DbLogMessage) -> Self {
        let level: JobLogLevel = match val.log_level {
            LogLevel::info => JobLogLevel::Info,
            LogLevel::warn => JobLogLevel::Warn,
            LogLevel::error => JobLogLevel::Error,
        };

        JobLogMessage {
            id: val.pub_id,
            created_at: to_micros(val.created_at),
            operator_id: val.operator_id,
            task_index: val.task_index.map(|i| i as u64),
            level,
            message: val.message,
            details: val.details,
            error_domain: ErrorDomain::from_str(&val.error_domain).ok(),
        }
    }
}

/// List a job's checkpoints
#[utoipa::path(
    get,
    path = "/v1/pipelines/{pipeline_id}/jobs/{job_id}/checkpoints",
    tag = "jobs",
    params(
        ("pipeline_id" = String, Path, description = "Pipeline id"),
        ("job_id" = String, Path, description = "Job id")
    ),
    responses(
        (status = 200, description = "Got job's checkpoints", body = CheckpointCollection),
    ),
)]
pub async fn get_job_checkpoints(
    State(state): State<AppState>,
    bearer_auth: BearerAuth,
    Path(PipelineJobPath {
        id: pipeline_pub_id,
        job_id: job_pub_id,
    }): Path<PipelineJobPath>,
) -> Result<Json<CheckpointCollection>, ErrorResp> {
    let db = state.database.client().await?;
    let auth_data = authenticate(&state.database, bearer_auth).await?;

    let job = api_queries::fetch_get_pipeline_job(
        &db,
        &auth_data.organization_id,
        &pipeline_pub_id,
        &job_pub_id,
    )
    .await?
    .into_iter()
    .next()
    .ok_or_else(|| not_found("Job"))?;

    let state_context: Option<StateContext> = job
        .state_context
        .map(serde_json::from_value)
        .transpose()
        .with_context(|| format!("converting state context for job {}", job_pub_id))
        .map_err(log_and_map)?;

    let checkpoints = if let Some(state_context) = state_context
        && let Some(leader) = state_context.leader
    {
        fetch_from_leader(leader, move |mut client, generation| async move {
            client
                .get_job_checkpoints(GetJobCheckpointsReq {
                    job_id: job.id,
                    generation,
                })
                .await
        })
        .await?
        .into_inner()
        .checkpoints
        .into_iter()
        .map(Checkpoint::from)
        .collect()
    } else {
        api_queries::fetch_get_job_checkpoints(&db, &job_pub_id, &auth_data.organization_id)
            .await
            .map_err(log_and_map)?
            .into_iter()
            .filter_map(|m| m.try_into().ok())
            .collect()
    };

    Ok(Json(CheckpointCollection { data: checkpoints }))
}

/// Get a checkpoint's details
#[utoipa::path(
    get,
    path = "/v1/pipelines/{pipeline_id}/jobs/{job_id}/checkpoints/{epoch}/operator_checkpoint_groups",
    tag = "jobs",
    params(
        ("pipeline_id" = String, Path, description = "Pipeline id"),
        ("job_id" = String, Path, description = "Job id"),
        ("epoch" = u32, Path, description = "Epoch")
    ),
    responses(
        (status = 200, description = "Got checkpoint's details", body = OperatorCheckpointGroupCollection),
    ),
)]
pub async fn get_checkpoint_details(
    State(state): State<AppState>,
    bearer_auth: BearerAuth,
    Path(PipelineJobCheckpointPath {
        id: pipeline_pub_id,
        job_id: job_pub_id,
        epoch,
    }): Path<PipelineJobCheckpointPath>,
) -> Result<Json<OperatorCheckpointGroupCollection>, ErrorResp> {
    let db = state.database.client().await?;
    let auth_data = authenticate(&state.database, bearer_auth).await?;

    let job = api_queries::fetch_get_pipeline_job(
        &db,
        &auth_data.organization_id,
        &pipeline_pub_id,
        &job_pub_id,
    )
    .await?
    .into_iter()
    .next()
    .ok_or_else(|| not_found("Job"))?;

    let state_context: Option<StateContext> = job
        .state_context
        .map(serde_json::from_value)
        .transpose()
        .with_context(|| format!("converting state context for job {}", job_pub_id))
        .map_err(log_and_map)?;

    let operators = if let Some(state_context) = state_context
        && let Some(leader) = state_context.leader
    {
        fetch_from_leader(leader, move |mut client, generation| async move {
            client
                .get_checkpoint_details(GetCheckpointDetailsReq {
                    job_id: job.id,
                    generation,
                    epoch: epoch as u64,
                })
                .await
        })
        .await?
        .into_inner()
        .operators
    } else {
        let checkpoint_details = api_queries::fetch_get_checkpoint_details(
            &db,
            &job_pub_id,
            &auth_data.organization_id,
            &(epoch as i32),
        )
        .await
        .map_err(log_and_map)?
        .into_iter()
        .next()
        .ok_or_else(|| {
            not_found(&format!(
                "Checkpoint with epoch {epoch} for job '{job_pub_id}'"
            ))
        })?;

        checkpoint_details
            .operators
            .map(|o| serde_json::from_value(o).unwrap())
            .unwrap_or_else(HashMap::<String, OperatorCheckpointDetail>::new)
    };

    let operators = operator_checkpoint_groups(operators);

    Ok(Json(OperatorCheckpointGroupCollection { data: operators }))
}

/// Subscribe to a job's output
#[utoipa::path(
    get,
    path = "/v1/pipelines/{pipeline_id}/jobs/{job_id}/output",
    tag = "jobs",
    params(
        ("pipeline_id" = String, Path, description = "Pipeline id"),
        ("job_id" = String, Path, description = "Job id")
    ),
    responses(
        (status = 200, description = "Job output as 'text/event-stream'"),
    ),
)]
pub async fn get_job_output(
    State(state): State<AppState>,
    bearer_auth: BearerAuth,
    Path(PipelineJobPath {
        id: pipeline_pub_id,
        job_id: job_pub_id,
    }): Path<PipelineJobPath>,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, ErrorResp> {
    let db = state.database.client().await?;
    let auth_data = authenticate(&state.database, bearer_auth).await?;

    // validate that the job exists, the user has access, and the graph has a GrpcSink
    query_job_by_pub_id(&pipeline_pub_id, &job_pub_id, &db, &auth_data).await?;

    let pipeline = query_pipeline_by_pub_id(&pipeline_pub_id, &db, &auth_data).await?;

    if !pipeline
        .graph
        .nodes
        .iter()
        .any(|n| n.operator.contains("preview"))
    {
        // TODO: make this check more robust
        return Err(bad_request("Job does not have a preview sink".to_string()));
    }
    let (tx, rx) = tokio::sync::mpsc::channel(32);

    let mut controller = controller_client("api", &config().api.tls)
        .await
        .map_err(log_and_map)?;

    let mut stream = controller
        .subscribe_to_output(Request::new(grpc::rpc::GrpcOutputSubscription {
            job_id: job_pub_id.clone(),
        }))
        .await
        .map_err(|e| match e.code() {
            Code::FailedPrecondition | Code::NotFound => bad_request(e.message().to_string()),
            _ => log_and_map(e),
        })?
        .into_inner();

    info!("Subscribed to output");
    tokio::spawn(async move {
        let _controller = controller;
        let mut message_count = 0;
        while let Some(d) = stream.next().await {
            if d.as_ref().map(|t| t.done).unwrap_or(false) {
                info!("Stream done for {}", job_pub_id);
                break;
            }

            let v = match d {
                Ok(d) => d,
                Err(_) => break,
            };

            let output_data: OutputData = v.into();

            let e = Ok(Event::default()
                .json_data(output_data)
                .unwrap()
                .id(message_count.to_string()));

            if tx.send(e).await.is_err() {
                break;
            }

            message_count += 1;
        }

        info!("Closing watch stream for {}", job_pub_id);
    });

    Ok(Sse::new(ReceiverStream::new(rx)))
}

/// Get all jobs
#[utoipa::path(
    get,
    path = "/v1/jobs",
    tag = "jobs",
    responses(
        (status = 200, description = "Get all jobs", body = JobCollection),
    ),
)]
pub async fn get_jobs(
    State(state): State<AppState>,
    bearer_auth: BearerAuth,
) -> Result<Json<JobCollection>, ErrorResp> {
    let auth_data = authenticate(&state.database, bearer_auth).await?;

    let jobs: Vec<DbPipelineJob> = api_queries::fetch_get_all_jobs(
        &state.database.client().await?,
        &auth_data.organization_id,
    )
    .await?;

    Ok(Json(JobCollection {
        data: jobs.into_iter().map(|p| p.into()).collect(),
    }))
}

impl TryFrom<DbCheckpoint> for Checkpoint {
    type Error = anyhow::Error;

    fn try_from(val: DbCheckpoint) -> anyhow::Result<Self> {
        let events: Vec<JobCheckpointSpan> = serde_json::from_value(val.event_spans)?;

        Ok(Checkpoint {
            epoch: val.epoch as u64,
            backend: val.state_backend,
            checkpoint_type: CheckpointType::from_is_stopping(val.is_stopping),
            start_time: to_micros(val.start_time),
            finish_time: val.finish_time.map(to_micros),
            events: events.into_iter().map(|e| e.into()).collect(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::OrgMetadata;
    use std::sync::{Arc, Mutex};

    fn auth() -> AuthData {
        AuthData {
            user_id: "user_2".to_string(),
            organization_id: "org_1".to_string(),
            role: "user".to_string(),
            org_metadata: OrgMetadata::default(),
        }
    }

    fn sqlite_job(state: &str) -> DatabaseSource {
        let connection = rusqlite::Connection::open_in_memory().unwrap();
        connection
            .execute_batch(
                "CREATE TABLE pipelines (
                    id INTEGER PRIMARY KEY,
                    pub_id TEXT NOT NULL,
                    organization_id TEXT NOT NULL
                );
                CREATE TABLE job_configs (
                    id TEXT PRIMARY KEY,
                    organization_id TEXT,
                    pipeline_name TEXT NOT NULL,
                    created_by TEXT NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL,
                    updated_by TEXT,
                    updated_at TIMESTAMP,
                    ttl_micros INTEGER,
                    stop TEXT DEFAULT 'none' NOT NULL,
                    parallelism_overrides TEXT DEFAULT '{}' NOT NULL,
                    checkpoint_interval_micros INTEGER DEFAULT 10000000 NOT NULL,
                    pipeline_id INTEGER NOT NULL,
                    restart_nonce INTEGER DEFAULT 0 NOT NULL,
                    restart_mode TEXT DEFAULT 'safe' NOT NULL,
                    ignore_state_before_epoch INTEGER,
                    env_vars TEXT DEFAULT '{}' NOT NULL,
                    scheduler_config TEXT DEFAULT '{}' NOT NULL
                );
                CREATE TABLE job_statuses (
                    pub_id TEXT NOT NULL UNIQUE,
                    id TEXT PRIMARY KEY,
                    organization_id TEXT,
                    run_id INTEGER DEFAULT 0 NOT NULL,
                    state TEXT DEFAULT 'Created' NOT NULL
                );
                CREATE TABLE checkpoints (job_id TEXT NOT NULL);
                CREATE TABLE job_log_messages (job_id TEXT NOT NULL);
                INSERT INTO pipelines (id, pub_id, organization_id)
                    VALUES (1, 'pl_1', 'org_1');
                INSERT INTO job_configs
                    (id, organization_id, pipeline_name, created_by, updated_by, ttl_micros, stop,
                     parallelism_overrides, checkpoint_interval_micros, pipeline_id, restart_nonce,
                     restart_mode, ignore_state_before_epoch, env_vars, scheduler_config)
                    VALUES
                    ('job_old', 'org_1', 'pipeline', 'user_1', 'user_1', 123, 'immediate',
                     '{\"1\": 4}', 5000000, 1, 3, 'force', 42,
                     '{\"ENV\": \"value\"}', '{\"scheduler\": true}');
                INSERT INTO checkpoints (job_id) VALUES ('job_old');
                INSERT INTO job_log_messages (job_id) VALUES ('job_old');",
            )
            .unwrap();
        connection
            .execute(
                "INSERT INTO job_statuses (pub_id, id, organization_id, run_id, state)
                 VALUES ('js_old', 'job_old', 'org_1', 7, ?1)",
                [state],
            )
            .unwrap();

        DatabaseSource::Sqlite(Arc::new(Mutex::new(connection)))
    }

    #[tokio::test]
    async fn replace_job_without_state_replaces_terminal_job_and_deletes_history() {
        let database = sqlite_job("Stopped");

        let new_job_id = replace_job_without_state(&database, "pl_1", &auth())
            .await
            .unwrap();

        assert_ne!(new_job_id, "job_old");
        let DatabaseSource::Sqlite(connection) = &database else {
            unreachable!()
        };
        let connection = connection.lock().unwrap();

        let job: (
            String,
            String,
            i64,
            String,
            String,
            Option<i64>,
            String,
            String,
        ) = connection
            .query_row(
                "SELECT id, stop, restart_nonce, restart_mode, updated_by,
                        ignore_state_before_epoch, env_vars, scheduler_config
                 FROM job_configs",
                [],
                |row| {
                    Ok((
                        row.get(0)?,
                        row.get(1)?,
                        row.get(2)?,
                        row.get(3)?,
                        row.get(4)?,
                        row.get(5)?,
                        row.get(6)?,
                        row.get(7)?,
                    ))
                },
            )
            .unwrap();
        assert_eq!(
            job,
            (
                new_job_id.clone(),
                "none".to_string(),
                0,
                "safe".to_string(),
                "user_2".to_string(),
                None,
                "{\"ENV\": \"value\"}".to_string(),
                "{\"scheduler\": true}".to_string(),
            )
        );

        let status: (String, i64, String) = connection
            .query_row("SELECT id, run_id, state FROM job_statuses", [], |row| {
                Ok((row.get(0)?, row.get(1)?, row.get(2)?))
            })
            .unwrap();
        assert_eq!(status, (new_job_id, 0, "Created".to_string()));

        for table in ["checkpoints", "job_log_messages"] {
            let count: i64 = connection
                .query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |row| {
                    row.get(0)
                })
                .unwrap();
            assert_eq!(count, 0, "{table} should be cleared");
        }
    }

    #[tokio::test]
    async fn replace_job_without_state_rejects_non_terminal_job() {
        let database = sqlite_job("Running");

        let error = replace_job_without_state(&database, "pl_1", &auth())
            .await
            .unwrap_err();

        assert_eq!(error.status_code, StatusCode::CONFLICT);
        let DatabaseSource::Sqlite(connection) = &database else {
            unreachable!()
        };
        let connection = connection.lock().unwrap();
        let job_id: String = connection
            .query_row("SELECT id FROM job_configs", [], |row| row.get(0))
            .unwrap();
        assert_eq!(job_id, "job_old");
    }
}