lix 0.12.0

Embeddable version control for apps and AI agents.
Documentation
use std::sync::Arc;

use async_trait::async_trait;
use datafusion::arrow::array::{Array, StringArray};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::{DataFusionError, Result};
use datafusion::execution::context::ExecutionProps;
use datafusion::logical_expr::Expr;
use datafusion::physical_plan::ExecutionPlan;

use crate::LixError;
use crate::sql2::error::lix_error_to_datafusion_error;
use crate::sql2::{DiffCommand, SqlWriteContext, WriteAccess};

use super::spec::{InsertApply, PlannedScan, TableSpec, register_spec_table, scan_row_source};

pub(super) async fn register_diff_command_provider(
    session: &datafusion::prelude::SessionContext,
    surface_name: &str,
    command: DiffCommand,
    write_ctx: SqlWriteContext,
) -> Result<(), LixError> {
    register_spec_table(
        session,
        surface_name,
        Arc::new(DiffCommandSpec {
            table_name: surface_name.to_string(),
            command,
        }),
        WriteAccess::write(write_ctx),
    )
}

struct DiffCommandSpec {
    table_name: String,
    command: DiffCommand,
}

#[async_trait]
impl TableSpec for DiffCommandSpec {
    fn table_name(&self) -> &str {
        &self.table_name
    }

    fn schema(&self) -> SchemaRef {
        command_schema()
    }

    async fn plan_scan(
        &self,
        _projection: Option<&Vec<usize>>,
        _filters: &[Expr],
        _limit: Option<usize>,
        _props: &ExecutionProps,
    ) -> Result<PlannedScan> {
        let table = self.table_name.clone();
        Ok(PlannedScan {
            schema: command_schema(),
            ordering: None,
            source: scan_row_source(command_schema(), table, |table| async move {
                Err(DataFusionError::Execution(format!(
                    "{table} is a command sink and cannot be read"
                )))
            }),
        })
    }

    async fn plan_insert(
        &self,
        write_ctx: SqlWriteContext,
        _input: &Arc<dyn ExecutionPlan>,
    ) -> Result<Option<InsertApply>> {
        let command = self.command;
        Ok(Some(Arc::new(move |batches| {
            let write_ctx = write_ctx.clone();
            Box::pin(async move {
                let diff_ids = diff_ids_from_batches(&batches)?;
                if diff_ids.is_empty() {
                    return Ok(0);
                }
                write_ctx
                    .execute_diff_command(command, diff_ids)
                    .await
                    .map(|outcome| outcome.rows_affected)
                    .map_err(lix_error_to_datafusion_error)
            })
        })))
    }
}

fn command_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("diff_id", DataType::Utf8, false),
        Field::new("commit_id", DataType::Utf8, false),
    ]))
}

fn diff_ids_from_batches(batches: &[RecordBatch]) -> Result<Vec<String>> {
    let mut ids = Vec::new();
    for batch in batches {
        let values = batch
            .column_by_name("diff_id")
            .ok_or_else(|| DataFusionError::Execution("diff_id column is required".to_string()))?
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or_else(|| DataFusionError::Execution("diff_id must be text".to_string()))?;
        for index in 0..values.len() {
            if values.is_null(index) {
                return Err(DataFusionError::Execution(
                    "diff_id cannot be NULL".to_string(),
                ));
            }
            ids.push(values.value(index).to_string());
        }
    }
    ids.sort();
    if ids.windows(2).any(|pair| pair[0] == pair[1]) {
        return Err(DataFusionError::Execution(
            "diff command selection contains duplicate diff_id rows".to_string(),
        ));
    }
    Ok(ids)
}