spiral-sdk 0.1.11

//! Enrichments: one named SPQL expression over each row of a source Table, computed by refresh
//! Tasks and kept as a mirror of the source's files.
//!
//! An Enrichment version records the expression, the output type it had when created, and the
//! source version its values are current for. Its values live in a private backing Table whose
//! files each mirror a source file: the same file ID, row count and row IDs. A refresh computes
//! only the rows whose inputs changed since that source version, and copies the values of rows
//! a rewrite such as compaction carried unchanged.

use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::api::enrichments::{Coverage, Definition, NativeState, Payload};
use crate::api::tasks::{
    EnrichmentRefreshDefinition, TaskDefinition, TaskReadTarget, TaskReadTargetCommitOnBranch,
};
use crate::catalog_types::{ProjectBranchId, ProjectCommitId};
use crate::{Error, Expression, Project, ResourceAddress, Run, Snapshot, TaskHandle};

pub(crate) const ENRICHMENT_KIND: &str = "enrichment";
const PRODUCER_CODEC: &str = "scandal.expression";
const OUTPUT_CODEC: &str = "arrow.field";
const COVERAGE_CODEC: &str = "spiral.enrichment-coverage";
const MIRROR_CODEC: &str = "scandal.mirror";
const DEFINITION_VERSION: u32 = 1;

/// The column of an Enrichment's backing Table that holds its values.
pub const ENRICHMENT_VALUE_COLUMN: &str = "value";

/// An Enrichment version's definition: its expression and the output type it had when created.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EnrichmentDefinition {
    pub expression: Expression,
    pub output: arrow_schema::Field,
}

#[derive(Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct Producer {
    expression: Expression,
}

#[derive(Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct EnrichmentCoverage {
    computed: bool,
}

#[derive(Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
struct Mirror {}

impl EnrichmentDefinition {
    /// The stored form of this definition.
    pub fn to_wire(&self) -> Result<Definition, Error> {
        Ok(Definition {
            schema_version: DEFINITION_VERSION,
            output_dtype: payload(OUTPUT_CODEC, &self.output)?,
            // A refresh derives the columns the expression reads from the expression itself.
            dependency_field_ids: Vec::new(),
            producer: payload(
                PRODUCER_CODEC,
                &Producer {
                    expression: self.expression.clone(),
                },
            )?,
        })
    }

    /// Decode a stored definition. Definitions that only the retired query engine can run, such
    /// as SQL expressions, are refused.
    pub fn from_wire(definition: &Definition) -> Result<Self, Error> {
        if definition.schema_version != DEFINITION_VERSION {
            return Err(Error::invalid(format!(
                "unsupported Enrichment definition version {}",
                definition.schema_version
            )));
        }
        let producer: Producer = decode(PRODUCER_CODEC, &definition.producer)?;
        Ok(Self {
            expression: producer.expression,
            output: decode(OUTPUT_CODEC, &definition.output_dtype)?,
        })
    }
}

/// Coverage of a version whose values a refresh computed, or of one created without values.
pub fn enrichment_coverage(computed: bool) -> Result<Coverage, Error> {
    Ok(Coverage {
        state: payload(COVERAGE_CODEC, &EnrichmentCoverage { computed })?,
    })
}

/// The native state that points an Enrichment version at its backing Table's version.
pub fn mirror_state(backing_version_id: Uuid) -> Result<NativeState, Error> {
    Ok(NativeState {
        backing_version_id,
        state: payload(MIRROR_CODEC, &Mirror {})?,
    })
}

/// The backing Table version a version's native state names, refusing state of other formats.
pub fn mirrored_version(state: &NativeState) -> Result<Uuid, Error> {
    let Mirror {} = decode(MIRROR_CODEC, &state.state)?;
    Ok(state.backing_version_id)
}

fn payload(codec: &str, value: &impl Serialize) -> Result<Payload, Error> {
    Ok(Payload {
        codec: codec.to_owned(),
        codec_version: 1,
        canonical_payload: serde_json::to_string(value).map_err(Error::invalid)?,
    })
}

fn decode<T: serde::de::DeserializeOwned>(codec: &str, payload: &Payload) -> Result<T, Error> {
    if payload.codec != codec || payload.codec_version != 1 {
        return Err(Error::invalid(format!(
            "unsupported Enrichment payload `{}` version {}, expected `{codec}` version 1",
            payload.codec, payload.codec_version
        )));
    }
    serde_json::from_str(&payload.canonical_payload).map_err(Error::invalid)
}

/// One Enrichment at the commit a [`Snapshot`] selects.
#[derive(Clone)]
pub struct Enrichment {
    project: Project,
    commit_id: ProjectCommitId,
    branch_id: Option<ProjectBranchId>,
    resource_id: Uuid,
    version_id: Uuid,
}

impl Snapshot {
    /// Resolve the Enrichment at `path`, such as `frames/is_blurry`, in this snapshot.
    pub async fn enrichment(&self, path: &str) -> Result<Enrichment, Error> {
        let address = path.parse::<ResourceAddress>()?;
        if address.reference().is_some() {
            return Err(Error::invalid(
                "select the Enrichment's branch with the snapshot, not an `@` reference",
            ));
        }
        let resource = self.resource_at_path(address.path()).await?;
        if resource.descriptor.kind.0 != ENRICHMENT_KIND {
            return Err(Error::invalid(format!("`{path}` is not an Enrichment")));
        }
        Ok(Enrichment {
            project: self.project.clone(),
            commit_id: self.commit_id,
            branch_id: self.branch_id,
            resource_id: resource.descriptor.resource_id.as_uuid(),
            version_id: resource.version_id.as_uuid(),
        })
    }
}

impl Enrichment {
    /// The Enrichment's stable resource identifier.
    pub fn resource_id(&self) -> Uuid {
        self.resource_id
    }

    /// The exact Enrichment version this handle resolved.
    pub fn version_id(&self) -> Uuid {
        self.version_id
    }

    /// Submit a refresh against the source version at this snapshot's commit.
    ///
    /// The first refresh computes every row; later ones compute only the rows whose inputs
    /// changed since the source version the values are current for, and copy the values of
    /// rows that maintenance moved. The refresh publishes one Enrichment version, which its
    /// result lists, or nothing when the values are already current. Returns `None` when a
    /// refresh of this Enrichment and branch is already waiting to run.
    pub async fn refresh(&self, run: Run) -> Result<Option<TaskHandle>, Error> {
        let branch_id = self.branch_id.ok_or_else(|| {
            Error::invalid("refreshing an Enrichment needs a branch snapshot, not an exact commit")
        })?;
        let definition = TaskDefinition::EnrichmentRefresh(EnrichmentRefreshDefinition {
            read_target: TaskReadTarget::CommitOnBranch(TaskReadTargetCommitOnBranch {
                branch_id: branch_id.as_uuid(),
                commit_id: self.commit_id.as_uuid(),
            }),
            enrichment_resource_id: self.resource_id,
        });
        self.project.submit_task(definition, None, run).await
    }
}