crank/crates/crank-registry/src/postgres/import_job.rs

use super::*;
use crate::{ArtifactSourceId, ImportJobSourceEnvelope};
use crank_artifacts::ArtifactRef;

const APPLICATION_RESULT_KEY: &str = "_crank_application_result";
const IMPORT_JOB_CLEANUP_BATCH: i64 = 128;
const DANGLING_OPENAPI_SOURCE_GRACE: time::Duration = time::Duration::minutes(5);

impl PostgresRegistry {
    pub async fn create_import_job(
        &self,
        request: CreateImportJobRequest<'_>,
    ) -> Result<(), RegistryError> {
        validate_source_envelope(request.source, request.preview_payload)?;
        let mut transaction = self.pool.begin().await?;
        sqlx::query(
            "insert into import_jobs (
                id,
                workspace_id,
                kind,
                source_format,
                source_version,
                status,
                preview_payload,
                created_operation_ids,
                error_text,
                created_at,
                expires_at,
                finished_at
             ) values (
                $1, $2, $3, $4, $5, $6, $7, '[]'::jsonb, null, $8::timestamptz, $9::timestamptz, null
             )",
        )
        .bind(request.id.as_str())
        .bind(request.workspace_id.as_str())
        .bind(serialize_enum_text(&request.kind, "kind")?)
        .bind(request.source_format)
        .bind(request.source_version)
        .bind(serialize_enum_text(&request.status, "status")?)
        .bind(request.preview_payload)
        .bind(request.created_at)
        .bind(request.expires_at)
        .execute(&mut *transaction)
        .await?;
        transaction.commit().await?;

        Ok(())
    }

    pub async fn get_import_job(
        &self,
        workspace_id: &WorkspaceId,
        job_id: &ImportJobId,
    ) -> Result<Option<ImportJob>, RegistryError> {
        let row = sqlx::query(
            "select
                id,
                workspace_id,
                kind,
                source_format,
                source_version,
                status,
                preview_payload,
                created_operation_ids,
                error_text,
                created_at,
                expires_at,
                finished_at
             from import_jobs
             where id = $1 and workspace_id = $2",
        )
        .bind(job_id.as_str())
        .bind(workspace_id.as_str())
        .fetch_optional(&self.pool)
        .await?;

        row.as_ref().map(map_import_job).transpose()
    }

    pub async fn finish_import_job(
        &self,
        request: FinishImportJobRequest<'_>,
    ) -> Result<(), RegistryError> {
        let result = sqlx::query(
            "update import_jobs
                set status = $2,
                    created_operation_ids = $3,
                    error_text = $4,
                    finished_at = $5::timestamptz
              where id = $1",
        )
        .bind(request.id.as_str())
        .bind(serialize_enum_text(&request.status, "status")?)
        .bind(request.created_operation_ids)
        .bind(request.error_text)
        .bind(request.finished_at)
        .execute(&self.pool)
        .await?;

        if result.rows_affected() == 0 {
            return Err(RegistryError::ImportJobNotFound {
                job_id: request.id.as_str().to_owned(),
            });
        }

        Ok(())
    }

    pub async fn apply_import_job(
        &self,
        request: ApplyImportJobRequest<'_>,
    ) -> Result<ImportJobApplyResult, RegistryError> {
        let mut tx = self.pool.begin().await?;
        let applied = apply_import_job_transaction(&mut tx, &request).await;

        match applied {
            Ok(result) => {
                tx.commit().await?;
                Ok(result)
            }
            Err(error) => {
                tx.rollback().await?;
                let error_text = error.to_string();
                let _ = sqlx::query(
                    "update import_jobs
                     set status = $3,
                         error_text = $4,
                         finished_at = $5::timestamptz
                     where id = $1
                       and workspace_id = $2
                       and status <> $6",
                )
                .bind(request.id.as_str())
                .bind(request.workspace_id.as_str())
                .bind(serialize_enum_text(&ImportJobStatus::Failed, "status")?)
                .bind(error_text)
                .bind(request.finished_at)
                .bind(serialize_enum_text(&ImportJobStatus::Completed, "status")?)
                .execute(&self.pool)
                .await;
                Err(error)
            }
        }
    }

    pub async fn delete_expired_import_jobs(&self) -> Result<u64, RegistryError> {
        let mut transaction = self.pool.begin().await?;
        let now = sqlx::query_scalar::<_, OffsetDateTime>("select now()")
            .fetch_one(&mut *transaction)
            .await?;
        let expired = sqlx::query(
            "select id, workspace_id, preview_payload
               from import_jobs
              where expires_at < now()
              order by expires_at, id
              limit $1
              for update skip locked",
        )
        .bind(IMPORT_JOB_CLEANUP_BATCH)
        .fetch_all(&mut *transaction)
        .await?;
        let mut expired_ids = Vec::with_capacity(expired.len());
        for row in &expired {
            expired_ids.push(row.try_get::<String, _>("id")?);
            let workspace_id = WorkspaceId::new(row.try_get::<String, _>("workspace_id")?);
            let payload = row.try_get::<Value, _>("preview_payload")?;
            if let Some(source) = source_from_payload(&payload)? {
                detach_source_in_transaction(
                    &mut transaction,
                    &workspace_id,
                    &source.source_id,
                    now,
                )
                .await?;
            }
        }
        let deleted = if expired_ids.is_empty() {
            0
        } else {
            sqlx::query("delete from import_jobs where id = any($1)")
                .bind(&expired_ids)
                .execute(&mut *transaction)
                .await?
                .rows_affected()
        };

        let dangling_cutoff = now.checked_sub(DANGLING_OPENAPI_SOURCE_GRACE).ok_or(
            RegistryError::InvalidArtifactSource {
                field: "import_job_source",
            },
        )?;
        let dangling = sqlx::query(
            "select s.workspace_id, s.source_id
               from artifact_sources s
              where s.lifecycle = 'active'
                and left(s.source_id, 12) = 'src_openapi_'
                and s.created_at < $1
                and not exists (
                    select 1
                      from import_jobs j
                     where j.workspace_id = s.workspace_id
                       and j.preview_payload -> 'source' ->> 'source_id' = s.source_id
                )
              order by s.created_at, s.workspace_id, s.source_id
              limit $2
              for update of s skip locked",
        )
        .bind(dangling_cutoff)
        .bind(IMPORT_JOB_CLEANUP_BATCH)
        .fetch_all(&mut *transaction)
        .await?;
        for row in dangling {
            let workspace_id = WorkspaceId::new(row.try_get::<String, _>("workspace_id")?);
            let source_id = ArtifactSourceId::new(row.try_get::<String, _>("source_id")?);
            detach_source_in_transaction(&mut transaction, &workspace_id, &source_id, now).await?;
        }
        transaction.commit().await?;

        Ok(deleted)
    }
}

async fn apply_import_job_transaction(
    tx: &mut Transaction<'_, Postgres>,
    request: &ApplyImportJobRequest<'_>,
) -> Result<ImportJobApplyResult, RegistryError> {
    let row = sqlx::query(
        "select status, preview_payload
         from import_jobs
         where id = $1 and workspace_id = $2
         for update",
    )
    .bind(request.id.as_str())
    .bind(request.workspace_id.as_str())
    .fetch_optional(&mut **tx)
    .await?
    .ok_or_else(|| RegistryError::ImportJobNotFound {
        job_id: request.id.as_str().to_owned(),
    })?;
    let status = deserialize_enum_text::<ImportJobStatus>(row.try_get("status")?, "status")?;
    let mut preview_payload = row.try_get::<Value, _>("preview_payload")?;
    let source =
        source_from_payload(&preview_payload)?.ok_or(RegistryError::InvalidArtifactSource {
            field: "import_job_source",
        })?;

    if status == ImportJobStatus::Completed
        && let Some(result) = stored_application_result(&preview_payload)?
    {
        if result.application_key != request.application_key {
            return Err(RegistryError::ImportJobAlreadyApplied {
                job_id: request.id.as_str().to_owned(),
            });
        }
        return Ok(result);
    }

    sqlx::query("select id from workspaces where id = $1 for update")
        .bind(request.workspace_id.as_str())
        .fetch_one(&mut **tx)
        .await?;

    let recorded_source = sqlx::query(
        "select s.lifecycle, b.artifact_ref
           from artifact_sources s
           join artifact_blobs b on b.digest = s.blob_digest
          where s.workspace_id = $1 and s.source_id = $2
          for update",
    )
    .bind(request.workspace_id.as_str())
    .bind(source.source_id.as_str())
    .fetch_optional(&mut **tx)
    .await?;
    let source_is_current = recorded_source
        .as_ref()
        .map(|row| {
            let lifecycle = row.try_get::<String, _>("lifecycle").ok();
            let artifact_ref = row.try_get::<String, _>("artifact_ref").ok();
            lifecycle.as_deref() == Some("active")
                && artifact_ref.as_deref() == Some(source.digest.as_str())
        })
        .unwrap_or(false);
    if !source_is_current {
        return Err(RegistryError::SourceUnavailable);
    }

    let mut result = ImportJobApplyResult {
        application_key: request.application_key.to_owned(),
        ..ImportJobApplyResult::default()
    };
    for draft in request.operations {
        if draft.operation.version != 1 {
            return Err(RegistryError::InvalidInitialVersion {
                operation_id: draft.operation.id.as_str().to_owned(),
                version: draft.operation.version,
            });
        }

        let mut operation = draft.operation.clone();
        let original_name = operation.name.clone();
        if operation_name_exists(tx, request.workspace_id, &operation.name).await? {
            match request.conflict_mode {
                ImportConflictMode::Skip => {
                    result.skipped.push(SkippedImportOperation {
                        operation_key: draft.operation_key.clone(),
                        name: operation.name,
                        reason: "operation_name_conflict".to_owned(),
                    });
                    continue;
                }
                ImportConflictMode::Rename => {
                    operation.name =
                        next_available_operation_name(tx, request.workspace_id, &operation.name)
                            .await?;
                }
            }
        }

        insert_operation_rows(tx, request.workspace_id, &operation, None).await?;
        result.created.push(AppliedImportOperation {
            operation_key: draft.operation_key.clone(),
            operation_id: operation.id,
            name: operation.name.clone(),
            version: operation.version,
            renamed_from: (operation.name != original_name).then_some(original_name),
        });
    }

    let stored_result = serde_json::to_value(&result)?;
    if let Some(object) = preview_payload.as_object_mut() {
        object.insert(APPLICATION_RESULT_KEY.to_owned(), stored_result);
    } else {
        preview_payload = serde_json::json!({
            "preview": preview_payload,
            "_crank_application_result": stored_result,
        });
    }
    let created_operation_ids = serde_json::to_value(
        result
            .created
            .iter()
            .map(|operation| operation.operation_id.as_str())
            .collect::<Vec<_>>(),
    )?;
    sqlx::query(
        "update import_jobs
         set status = $3,
             preview_payload = $4,
             created_operation_ids = $5,
             error_text = null,
             finished_at = $6::timestamptz
         where id = $1 and workspace_id = $2",
    )
    .bind(request.id.as_str())
    .bind(request.workspace_id.as_str())
    .bind(serialize_enum_text(&ImportJobStatus::Completed, "status")?)
    .bind(preview_payload)
    .bind(created_operation_ids)
    .bind(request.finished_at)
    .execute(&mut **tx)
    .await?;

    detach_source_in_transaction(
        tx,
        request.workspace_id,
        &source.source_id,
        *request.finished_at,
    )
    .await?;

    Ok(result)
}

fn validate_source_envelope(
    source: &ImportJobSourceEnvelope,
    preview_payload: &Value,
) -> Result<(), RegistryError> {
    if source.source_id.as_str().len() > 132 {
        return Err(RegistryError::InvalidArtifactSource {
            field: "import_job_source",
        });
    }
    let parsed =
        source_from_payload(preview_payload)?.ok_or(RegistryError::InvalidArtifactSource {
            field: "import_job_source",
        })?;
    if parsed != *source {
        return Err(RegistryError::InvalidArtifactSource {
            field: "import_job_source",
        });
    }
    Ok(())
}

fn source_from_payload(payload: &Value) -> Result<Option<ImportJobSourceEnvelope>, RegistryError> {
    let Some(source) = payload.get("source") else {
        return Ok(None);
    };
    let has_source_id = source.get("source_id").is_some();
    let has_digest = source.get("digest").is_some();
    if !has_source_id && !has_digest {
        // Jobs created before the verified-source envelope used `source` for
        // OpenAPI format/version metadata. Expiry cleanup must remain able to
        // delete those jobs during a rolling upgrade.
        return Ok(None);
    }
    let source_id = source
        .get("source_id")
        .and_then(Value::as_str)
        .filter(|value| value.len() <= 132)
        .ok_or(RegistryError::InvalidArtifactSource {
            field: "import_job_source",
        })?;
    let digest = source.get("digest").and_then(Value::as_str).ok_or(
        RegistryError::InvalidArtifactSource {
            field: "import_job_source",
        },
    )?;
    let digest = ArtifactRef::parse(digest).map_err(|_| RegistryError::InvalidArtifactSource {
        field: "import_job_source",
    })?;
    Ok(Some(ImportJobSourceEnvelope {
        source_id: ArtifactSourceId::new(source_id),
        digest,
    }))
}

async fn detach_source_in_transaction(
    transaction: &mut Transaction<'_, Postgres>,
    workspace_id: &WorkspaceId,
    source_id: &ArtifactSourceId,
    detached_at: OffsetDateTime,
) -> Result<(), RegistryError> {
    sqlx::query(
        "update artifact_sources
            set lifecycle = 'detached', updated_at = $1, detached_at = $1
          where workspace_id = $2 and source_id = $3
            and lifecycle = 'active'",
    )
    .bind(detached_at)
    .bind(workspace_id.as_str())
    .bind(source_id.as_str())
    .execute(&mut **transaction)
    .await?;
    Ok(())
}

fn stored_application_result(
    preview_payload: &Value,
) -> Result<Option<ImportJobApplyResult>, RegistryError> {
    preview_payload
        .get(APPLICATION_RESULT_KEY)
        .cloned()
        .map(serde_json::from_value)
        .transpose()
        .map_err(RegistryError::from)
}

async fn operation_name_exists(
    tx: &mut Transaction<'_, Postgres>,
    workspace_id: &WorkspaceId,
    name: &str,
) -> Result<bool, RegistryError> {
    Ok(sqlx::query_scalar::<_, bool>(
        "select exists(
            select 1 from operations where workspace_id = $1 and name = $2
         )",
    )
    .bind(workspace_id.as_str())
    .bind(name)
    .fetch_one(&mut **tx)
    .await?)
}

async fn next_available_operation_name(
    tx: &mut Transaction<'_, Postgres>,
    workspace_id: &WorkspaceId,
    base_name: &str,
) -> Result<String, RegistryError> {
    for index in 2.. {
        let candidate = format!("{base_name}_{index}");
        if !operation_name_exists(tx, workspace_id, &candidate).await? {
            return Ok(candidate);
        }
    }

    unreachable!()
}