spiral-sdk 0.1.12

//! The Arrow framing of the scandal API, shared by every transport. The scandal contract
//! (`spiraldb/services/scandal.yaml`) defines the routes and JSON requests; rows travel as
//! Arrow IPC streams, and a `ReadScan` batch carries its progress in two trailing columns.

use std::sync::Arc;

use arrow_array::types::Int32Type;
use arrow_array::{
    Array as _, DictionaryArray, Int32Array, RecordBatch, RecordBatchOptions, StringArray,
    UInt64Array,
};
use arrow_ipc::writer::StreamWriter;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use bytes::Bytes;
use futures::stream::{self, BoxStream};
use futures::{StreamExt as _, TryStreamExt as _};

use super::error::ScandalError;
use super::plan::{PLAN_METADATA_KEY, PlanDescription};
use super::profile::{PROFILE_METADATA_KEY, execution_error};
use super::types::ScanMetricsSnapshot;
use super::{RecordBatchStream, SplitBatch};

/// Liveness. Not `/healthz`: Cloud Run reserves paths ending in `z` and never forwards them.
pub const HEALTH: &str = "/health";

/// The audience of a scandal capability: the RFC 8707 resource for which a client exchanges
/// its access token, and whose result every request except [`HEALTH`] carries as its bearer
/// token.
pub const CAPABILITY_AUDIENCE: &str = "https://scandal.spiraldb.dev";

/// The Arrow IPC streaming format's media type.
pub const ARROW_STREAM: &str = "application/vnd.apache.arrow.stream";
/// The column after each `ReadScan` batch's rows holding the next stored row it reached.
pub const NEXT_ROW_COLUMN: &str = "scandal.next_row";
/// The column after [`NEXT_ROW_COLUMN`] holding the counters that a split read accrued since
/// its previous batch.
pub const METRICS_COLUMN: &str = "scandal.metrics";

/// Encode batches as an Arrow IPC stream. Only a stream that ends without error receives
/// the end marker, so a reader detects truncation. Profiled failures end with an explicit
/// failed execution frame and normal IPC completion; adapters return its declared error.
pub fn encode_ipc(
    schema: SchemaRef,
    batches: RecordBatchStream,
) -> BoxStream<'static, Result<Bytes, ScandalError>> {
    let schema = match execution_schema(&schema) {
        Ok(schema) => schema,
        Err(error) => return stream::once(async move { Err(error) }).boxed(),
    };
    let framed_schema = Arc::clone(&schema);
    let profiled = schema.metadata().contains_key(EXECUTION_FRAME);
    let batches = batches
        .map(move |batch| {
            batch.and_then(|batch| {
                let failed = profiled
                    && batch
                        .schema_ref()
                        .metadata()
                        .contains_key(PROFILE_METADATA_KEY)
                    && PlanDescription::from_schema(batch.schema_ref())?
                        .as_ref()
                        .is_some_and(|p| execution_error(p).is_some());
                attach_execution(batch, &framed_schema).map(|batch| (batch, failed))
            })
        })
        .boxed();
    let writer = match StreamWriter::try_new(Vec::new(), &schema) {
        Ok(writer) => writer,
        Err(error) => return stream::once(async move { Err(internal(error)) }).boxed(),
    };
    stream::try_unfold(Some((batches, writer, true)), |state| async move {
        let Some((mut batches, mut writer, first)) = state else {
            return Ok(None);
        };
        if first {
            let header = Bytes::from(std::mem::take(writer.get_mut()));
            return Ok(Some((header, Some((batches, writer, false)))));
        }
        match batches.try_next().await? {
            Some((batch, failed)) => {
                writer.write(&batch).map_err(internal)?;
                if failed {
                    writer.finish().map_err(internal)?;
                    return Ok(Some((Bytes::from(std::mem::take(writer.get_mut())), None)));
                }
                let bytes = Bytes::from(std::mem::take(writer.get_mut()));
                Ok(Some((bytes, Some((batches, writer, false)))))
            }
            None => {
                writer.finish().map_err(internal)?;
                Ok(Some((Bytes::from(std::mem::take(writer.get_mut())), None)))
            }
        }
    })
    .boxed()
}

/// Encode application input rows, preserving execution options while removing observations
/// belonging to an earlier operation. Result framing must never become a stored input column.
pub fn encode_input(
    schema: SchemaRef,
    batches: RecordBatchStream,
) -> BoxStream<'static, Result<Bytes, ScandalError>> {
    let mut metadata = schema.metadata().clone();
    metadata.remove(PLAN_METADATA_KEY);
    metadata.remove(PROFILE_METADATA_KEY);
    encode_ipc(
        Arc::new(schema.as_ref().clone().with_metadata(metadata)),
        batches,
    )
}

/// The schema of a `ReadScan` response: the split's rows, then [`NEXT_ROW_COLUMN`] and
/// [`METRICS_COLUMN`].
///
/// Each trailing column is a dictionary of at most one value that every row's key names, so
/// the value travels with its batch even when the batch holds no rows.
pub fn read_scan_schema(rows: &Schema) -> SchemaRef {
    let dictionary = |values| DataType::Dictionary(Box::new(DataType::Int32), Box::new(values));
    let mut fields = rows.fields().iter().cloned().collect::<Vec<_>>();
    fields.push(Arc::new(Field::new(
        NEXT_ROW_COLUMN,
        dictionary(DataType::UInt64),
        false,
    )));
    fields.push(Arc::new(Field::new(
        METRICS_COLUMN,
        dictionary(DataType::Utf8),
        true,
    )));
    Arc::new(Schema::new_with_metadata(fields, rows.metadata().clone()))
}

/// Append a split batch's next stored row and counters as single-value dictionaries.
pub fn attach_progress(batch: SplitBatch, schema: &SchemaRef) -> Result<RecordBatch, ScandalError> {
    let rows = batch.batch.num_rows();
    let next_row = DictionaryArray::<Int32Type>::try_new(
        Int32Array::from(vec![0; rows]),
        Arc::new(UInt64Array::from(vec![batch.next_row])),
    )
    .map_err(internal)?;
    let metrics = match batch.metrics {
        Some(metrics) => DictionaryArray::<Int32Type>::try_new(
            Int32Array::from(vec![0; rows]),
            Arc::new(StringArray::from(vec![
                serde_json::to_string(&metrics).map_err(internal)?,
            ])),
        ),
        None => DictionaryArray::<Int32Type>::try_new(
            Int32Array::new_null(rows),
            Arc::new(StringArray::from(Vec::<String>::new())),
        ),
    }
    .map_err(internal)?;
    let mut columns = batch.batch.columns().to_vec();
    columns.push(Arc::new(next_row));
    columns.push(Arc::new(metrics));
    RecordBatch::try_new_with_options(
        progress_execution_schema(&batch.batch, schema),
        columns,
        &RecordBatchOptions::new().with_row_count(Some(rows)),
    )
    .map_err(internal)
}

/// Separate a `ReadScan` batch's rows from its next stored row and counters.
pub fn detach_progress(batch: RecordBatch) -> Result<SplitBatch, ScandalError> {
    let rows = trailing_columns(batch.schema_ref())?;
    let value = |index: usize| {
        batch
            .column(index)
            .as_any()
            .downcast_ref::<DictionaryArray<Int32Type>>()
            .map(|dictionary| Arc::clone(dictionary.values()))
    };
    let next_row = value(rows)
        .and_then(|values| values.as_any().downcast_ref::<UInt64Array>().cloned())
        .filter(|values| values.len() == 1 && values.null_count() == 0)
        .map(|values| values.value(0))
        .ok_or_else(|| internal("scandal batch's next stored row is malformed"))?;
    let metrics = value(rows + 1)
        .and_then(|values| values.as_any().downcast_ref::<StringArray>().cloned())
        .filter(|values| values.len() <= 1 && values.null_count() == 0)
        .ok_or_else(|| internal("scandal batch's counters are malformed"))?;
    let metrics = metrics
        .iter()
        .flatten()
        .next()
        .map(serde_json::from_str::<ScanMetricsSnapshot>)
        .transpose()
        .map_err(internal)?;
    let data = batch
        .project(&(0..rows).collect::<Vec<_>>())
        .map_err(internal)?;
    Ok(SplitBatch {
        batch: data,
        next_row,
        metrics,
    })
}

/// The row schema of a `ReadScan` response, without its trailing progress columns.
pub fn read_scan_rows(schema: &Schema) -> Result<SchemaRef, ScandalError> {
    let rows = trailing_columns(schema)?;
    Ok(Arc::new(Schema::new_with_metadata(
        schema.fields()[..rows].to_vec(),
        schema.metadata().clone(),
    )))
}

/// The number of row columns before [`NEXT_ROW_COLUMN`] and [`METRICS_COLUMN`].
fn trailing_columns(schema: &Schema) -> Result<usize, ScandalError> {
    schema
        .fields()
        .len()
        .checked_sub(2)
        .filter(|rows| {
            schema.field(*rows).name() == NEXT_ROW_COLUMN
                && schema.field(rows + 1).name() == METRICS_COLUMN
        })
        .ok_or_else(|| internal("scandal stream lacks its progress columns"))
}

/// Optional execution evidence travels as a single dictionary value even on empty batches.
/// Arrow IPC and C Stream schemas are fixed, so changing only schema metadata loses updates.
pub const EXECUTION_COLUMN: &str = "scandal.execution";
const EXECUTION_FRAME: &str = "spiral.execution.frame";

/// The fixed transport schema for an operation that opted into profiling.
pub fn execution_schema(rows: &Schema) -> Result<SchemaRef, ScandalError> {
    if PlanDescription::from_schema(rows)?.is_none_or(|plan| plan.profile.is_none()) {
        return Ok(Arc::new(rows.clone()));
    }
    if rows.metadata().contains_key(EXECUTION_FRAME) {
        return Ok(Arc::new(rows.clone()));
    }
    if rows
        .fields()
        .iter()
        .any(|field| field.name() == EXECUTION_COLUMN)
    {
        return Err(internal("reserved execution column in result"));
    }
    let mut fields = rows.fields().to_vec();
    fields.push(Arc::new(Field::new(
        EXECUTION_COLUMN,
        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
        false,
    )));
    let mut metadata = rows.metadata().clone();
    metadata.insert(EXECUTION_FRAME.into(), "1".into());
    Ok(Arc::new(Schema::new_with_metadata(fields, metadata)))
}

/// Append the current plan/profile without repeating its JSON for every row.
pub fn attach_execution(
    batch: RecordBatch,
    schema: &SchemaRef,
) -> Result<RecordBatch, ScandalError> {
    if !schema.metadata().contains_key(EXECUTION_FRAME) {
        return Ok(batch);
    }
    let encoded = batch
        .schema_ref()
        .metadata()
        .get(PLAN_METADATA_KEY)
        .ok_or_else(|| internal("profiled batch lacks execution description"))?
        .clone();
    if encoded.len() > 4 * 1024 * 1024 {
        return Err(internal("execution description exceeds 4 MiB"));
    }
    let evidence = DictionaryArray::<Int32Type>::try_new(
        Int32Array::from(vec![0; batch.num_rows()]),
        Arc::new(StringArray::from(vec![encoded])),
    )
    .map_err(internal)?;
    let mut columns = batch.columns().to_vec();
    columns.push(Arc::new(evidence));
    RecordBatch::try_new_with_options(
        Arc::clone(schema),
        columns,
        &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
    )
    .map_err(internal)
}

/// Remove the execution transport field from a response schema.
pub fn execution_rows(schema: &Schema) -> Result<SchemaRef, ScandalError> {
    if !schema.metadata().contains_key(EXECUTION_FRAME) {
        return Ok(Arc::new(schema.clone()));
    }
    let count = schema
        .fields()
        .len()
        .checked_sub(1)
        .filter(|last| schema.field(*last).name() == EXECUTION_COLUMN)
        .ok_or_else(|| internal("execution frame is malformed"))?;
    let mut metadata = schema.metadata().clone();
    metadata.remove(EXECUTION_FRAME);
    Ok(Arc::new(Schema::new_with_metadata(
        schema.fields()[..count].to_vec(),
        metadata,
    )))
}

/// Restore an execution snapshot from fixed-schema Arrow IPC or C Stream delivery.
pub fn detach_execution(batch: RecordBatch) -> Result<RecordBatch, ScandalError> {
    if !batch.schema_ref().metadata().contains_key(EXECUTION_FRAME) {
        return Ok(batch);
    }
    let schema = execution_rows(batch.schema_ref())?;
    let values = batch
        .column(batch.num_columns() - 1)
        .as_any()
        .downcast_ref::<DictionaryArray<Int32Type>>()
        .and_then(|dictionary| dictionary.values().as_any().downcast_ref::<StringArray>())
        .filter(|values| values.len() == 1 && values.null_count() == 0)
        .ok_or_else(|| internal("execution evidence is malformed"))?;
    let mut metadata = schema.metadata().clone();
    metadata.insert(PLAN_METADATA_KEY.into(), values.value(0).into());
    let schema = schema.as_ref().clone().with_metadata(metadata);
    PlanDescription::from_schema(&schema)?;
    RecordBatch::try_new_with_options(
        Arc::new(schema),
        batch.columns()[..batch.num_columns() - 1].to_vec(),
        &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
    )
    .map_err(internal)
}

/// Keep a terminal failed report visible before returning its sanitized contract error.
/// This makes profiled failures independent of HTTP body truncation behavior.
pub fn execution_batches(input: RecordBatchStream) -> RecordBatchStream {
    stream::unfold(Some((input, None)), |state| async move {
        let (mut input, pending) = state?;
        if let Some(error) = pending {
            return Some((Err(error), None));
        }
        match input.next().await {
            Some(Ok(batch)) => {
                let error = match batch
                    .schema_ref()
                    .metadata()
                    .contains_key(PROFILE_METADATA_KEY)
                    .then(|| PlanDescription::from_schema(batch.schema_ref()))
                    .transpose()
                {
                    Ok(plan) => plan.flatten().as_ref().and_then(execution_error),
                    Err(error) => return Some((Err(error), None)),
                };
                Some((Ok(batch), Some((input, error))))
            }
            Some(Err(error)) => Some((Err(error), None)),
            None => None,
        }
    })
    .boxed()
}

fn progress_execution_schema(batch: &RecordBatch, schema: &SchemaRef) -> SchemaRef {
    let mut metadata = schema.metadata().clone();
    if let Some(plan) = batch.schema_ref().metadata().get(PLAN_METADATA_KEY) {
        metadata.insert(PLAN_METADATA_KEY.into(), plan.clone());
    }
    Arc::new(schema.as_ref().clone().with_metadata(metadata))
}

fn internal(message: impl ToString) -> ScandalError {
    ScandalError::Internal {
        message: message.to_string(),
    }
}

#[cfg(test)]
mod tests {
    use arrow_array::Int64Array;
    use spiral_contract::safelong::SafeUnsignedLong;

    use super::super::plan::PLAN_METADATA_KEY;
    use super::*;

    #[test]
    fn progress_round_trips_beside_split_rows() {
        let rows = Arc::new(Schema::new_with_metadata(
            vec![Field::new("value", DataType::Int64, false)],
            [
                ("application".into(), "retained".into()),
                (PLAN_METADATA_KEY.into(), "execution description".into()),
            ]
            .into(),
        ));
        let schema = read_scan_schema(&rows);
        let batch =
            RecordBatch::try_new(rows, vec![Arc::new(Int64Array::from(vec![4, 5]))]).unwrap();
        let metrics = ScanMetricsSnapshot {
            output_rows: SafeUnsignedLong(2),
            ..Default::default()
        };
        for (batch, metrics) in [(batch.clone(), None), (batch.slice(0, 0), Some(metrics))] {
            let attached = attach_progress(
                SplitBatch {
                    batch: batch.clone(),
                    next_row: u64::MAX,
                    metrics: metrics.clone(),
                },
                &schema,
            )
            .unwrap();
            let detached = detach_progress(attached).unwrap();
            assert_eq!(
                (detached.batch, detached.next_row, detached.metrics),
                (batch, u64::MAX, metrics)
            );
        }
        assert_eq!(
            read_scan_rows(&schema).unwrap().fields().len(),
            1,
            "progress columns are not row columns"
        );
    }
}