spiral-sdk 0.1.11

//! Table edits driven end to end through the SDK against an in-process scandal that holds one
//! fragment and records every write request.

use std::ops::Range;
use std::sync::Arc;

use arrow_array::cast::AsArray as _;
use arrow_array::types::{Int64Type, UInt64Type};
use arrow_array::{Array, BooleanArray, Int64Array, RecordBatch, StringArray, UInt64Array};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use arrow_select::filter::filter_record_batch;
use async_trait::async_trait;
use futures::{StreamExt as _, TryStreamExt as _, stream};
use parking_lot::Mutex;
use reqwest::Url;
use serde_json::json;
use uuid::Uuid;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};

use crate::scandal::api::{ArrowSchema, EditBase, ScanInfo as WireScanInfo};
use crate::scandal::blob::{Blob, BlobPart, BlobStat, BlobUpload};
use crate::scandal::error::ScandalError;
use crate::scandal::expression::NamedExpression;
use crate::scandal::functions::FunctionInfo;
use crate::scandal::types::{
    FragmentReceipt, PlannedSplit, PreparedWriteReceipt, RowRange, Scan as ScandalScan, ScanPlan,
    Split, SplitTake,
};
use crate::scandal::{
    ByteStream, ChangePlan, DeleteRowsRequest, Evaluation, PatchRowsRequest, PlanChangesRequest,
    RecordBatchStream, Scandal, SplitBatch, SplitRead, StorageBinding, TableBinding, WriteBase,
};
use crate::{
    BySource, CHANGE_KIND_COLUMN, ChangeImages, ChangeOptions, Client, ColumnPath, Matched,
    MergeActions, NotMatched, ROW_ID_COLUMN, SchemaChange, Table, TableIdentity,
};

/// What the scandal was asked to write.
#[derive(Default)]
struct Written {
    patches: Vec<RecordBatch>,
    deletions: Vec<RecordBatch>,
    appends: Vec<RecordBatch>,
    finalized: Vec<(WriteBase, Vec<FragmentReceipt>)>,
}

/// One fragment of rows `id`, `label` with row IDs, read whole whatever the scan's filter.
struct OneFragment {
    rows: RecordBatch,
    row_ids: Vec<u64>,
    fragment_id: Uuid,
    written: Mutex<Written>,
    next_version: Uuid,
}

impl OneFragment {
    fn new(next_version: Uuid) -> Self {
        Self {
            rows: rows(vec![1, 2, 3], vec!["a", "b", "c"]),
            row_ids: vec![10, 11, 12],
            fragment_id: Uuid::now_v7(),
            written: Mutex::default(),
            next_version,
        }
    }

    /// The rows `scan` returns: those with its row IDs, with its columns, then their row IDs
    /// when it asks for them.
    fn read(&self, scan: &ScandalScan) -> RecordBatch {
        let all = self.columns(scan);
        let Some(ranges) = &scan.row_ids else {
            return all;
        };
        let keep = self
            .row_ids
            .iter()
            .map(|id| {
                Some(
                    ranges
                        .iter()
                        .any(|range| (range.start.0..range.end.0).contains(id)),
                )
            })
            .collect::<BooleanArray>();
        filter_record_batch(&all, &keep).unwrap()
    }

    fn columns(&self, scan: &ScandalScan) -> RecordBatch {
        let schema = self.rows.schema();
        let names = scan.columns.clone().unwrap_or_else(|| {
            schema
                .fields()
                .iter()
                .map(|field| field.name().clone())
                .collect()
        });
        let mut fields = Vec::new();
        let mut columns = Vec::new();
        for name in &names {
            fields.push(schema.field_with_name(name).unwrap().clone());
            columns.push(Arc::clone(self.rows.column_by_name(name).unwrap()));
        }
        if let Some(row_ids) = &scan.row_id_column {
            fields.push(Field::new(row_ids, DataType::UInt64, false));
            columns.push(Arc::new(UInt64Array::from(self.row_ids.clone())));
        }
        RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
    }
}

fn schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, true),
        Field::new("label", DataType::Utf8, true),
    ]))
}

fn rows(ids: Vec<i64>, labels: Vec<&str>) -> RecordBatch {
    RecordBatch::try_new(
        schema(),
        vec![
            Arc::new(Int64Array::from(ids)),
            Arc::new(StringArray::from(labels)),
        ],
    )
    .unwrap()
}

fn unsupported() -> ScandalError {
    ScandalError::InvalidArgument {
        message: "the one-fragment scandal does not do this".into(),
    }
}

async fn collected(input: RecordBatchStream) -> Result<Vec<RecordBatch>, ScandalError> {
    input.try_collect().await
}

fn receipt(token: &str) -> Vec<FragmentReceipt> {
    vec![FragmentReceipt {
        token: token.to_owned(),
    }]
}

#[async_trait]
impl Scandal for OneFragment {
    async fn list_functions(&self) -> Result<Vec<FunctionInfo>, ScandalError> {
        Err(unsupported())
    }

    async fn plan_scan(
        &self,
        scan: &ScandalScan,
        _: &StorageBinding,
    ) -> Result<ScanPlan, ScandalError> {
        let rows = self.rows.num_rows() as u64;
        Ok(ScanPlan {
            info: WireScanInfo {
                source_fields: None,
                filter_clauses: Vec::new(),
                schema: ArrowSchema::from_schema(&self.read(scan).schema())?,
                statistics: None,
                write_layout: None,
            },
            metrics: None,
            splits: vec![PlannedSplit {
                split: Split {
                    fragment: 0.into(),
                    rows: RowRange::new(0..rows),
                },
                estimated_rows: Some(rows.into()),
                estimated_bytes: None,
                ordering: None,
                fragment_id: Some(self.fragment_id),
            }],
        })
    }

    async fn read_scan(
        &self,
        scan: &ScandalScan,
        split: &Split,
        _: &StorageBinding,
    ) -> Result<SplitRead, ScandalError> {
        let batch = self.read(scan);
        Ok(SplitRead {
            schema: batch.schema(),
            batches: stream::iter([Ok(SplitBatch {
                batch,
                next_row: split.rows.end.0,
                metrics: None,
            })])
            .boxed(),
        })
    }

    async fn take_rows(
        &self,
        _: &ScandalScan,
        _: &[SplitTake],
        _: &StorageBinding,
    ) -> Result<SplitRead, ScandalError> {
        Err(unsupported())
    }

    async fn read_blob(
        &self,
        _: &Blob,
        _: Option<Range<u64>>,
        _: &StorageBinding,
    ) -> Result<ByteStream, ScandalError> {
        Err(unsupported())
    }

    async fn stat_blob(&self, _: &Blob, _: &StorageBinding) -> Result<BlobStat, ScandalError> {
        Err(unsupported())
    }

    async fn write_blob(
        &self,
        _: &TableBinding,
        _: &str,
        _: u64,
        _: ByteStream,
    ) -> Result<Blob, ScandalError> {
        Err(unsupported())
    }

    async fn pack_blobs(
        &self,
        _: &TableBinding,
        _: Option<Uuid>,
        _: SchemaRef,
        _: RecordBatchStream,
    ) -> Result<Vec<RecordBatch>, ScandalError> {
        Err(unsupported())
    }

    async fn start_blob_upload(
        &self,
        _: &TableBinding,
        _: Uuid,
        _: &str,
        _: Option<u64>,
    ) -> Result<BlobUpload, ScandalError> {
        Err(unsupported())
    }

    async fn upload_blob_part(
        &self,
        _: &BlobUpload,
        _: u32,
        _: ByteStream,
        _: &StorageBinding,
    ) -> Result<BlobPart, ScandalError> {
        Err(unsupported())
    }

    async fn complete_blob_upload(
        &self,
        _: &BlobUpload,
        _: &[BlobPart],
        _: &StorageBinding,
    ) -> Result<Blob, ScandalError> {
        Err(unsupported())
    }

    async fn abort_blob_upload(
        &self,
        _: &BlobUpload,
        _: &StorageBinding,
    ) -> Result<(), ScandalError> {
        Err(unsupported())
    }

    async fn append_rows(
        &self,
        _: &TableBinding,
        _: Option<Uuid>,
        _: Option<Uuid>,
        _: SchemaRef,
        input: RecordBatchStream,
    ) -> Result<Vec<FragmentReceipt>, ScandalError> {
        let batches = collected(input).await?;
        self.written.lock().appends.extend(batches);
        Ok(receipt("append"))
    }

    async fn patch_rows(
        &self,
        _: &PatchRowsRequest,
        _: SchemaRef,
        input: RecordBatchStream,
    ) -> Result<Vec<FragmentReceipt>, ScandalError> {
        let batches = collected(input).await?;
        self.written.lock().patches.extend(batches);
        Ok(receipt("patch"))
    }

    async fn delete_rows(
        &self,
        _: &DeleteRowsRequest,
        _: SchemaRef,
        input: RecordBatchStream,
    ) -> Result<Vec<FragmentReceipt>, ScandalError> {
        let batches = collected(input).await?;
        self.written.lock().deletions.extend(batches);
        Ok(receipt("delete"))
    }

    /// Row 10 deleted, row 11 patched, row 12 appended.
    async fn plan_changes(&self, _: &PlanChangesRequest) -> Result<ChangePlan, ScandalError> {
        let range = |start: u64| json!([{"start": start, "end": start + 1}]);
        Ok(serde_json::from_value(json!({
            "fragments": [
                {"fragmentId": self.fragment_id, "rowCount": 3, "change": {
                    "type": "retained",
                    "retained": {"patched": range(11), "deleted": range(10)}
                }},
                {"fragmentId": Uuid::now_v7(), "rowCount": 1, "change": {
                    "type": "added",
                    "added": {"appended": range(12), "carried": [], "patched": []}
                }}
            ],
            "removedFragmentIds": [],
            "retypedColumns": []
        }))
        .unwrap())
    }

    async fn finalize_write(
        &self,
        table: &TableBinding,
        base: &WriteBase,
        receipts: &[FragmentReceipt],
    ) -> Result<PreparedWriteReceipt, ScandalError> {
        self.written
            .lock()
            .finalized
            .push((base.clone(), receipts.to_vec()));
        Ok(PreparedWriteReceipt {
            base_version: base.base_version()?,
            resource_id: table.table_id,
            sequence: 2.into(),
            version_id: self.next_version,
        })
    }

    async fn bind_projection(
        &self,
        _: &Schema,
        _: &[NamedExpression],
    ) -> Result<SchemaRef, ScandalError> {
        Err(unsupported())
    }

    async fn evaluate(
        &self,
        _: &[NamedExpression],
        _: SchemaRef,
        _: RecordBatchStream,
        _: &[TableBinding],
    ) -> Result<Evaluation, ScandalError> {
        Err(unsupported())
    }
}

/// A table of [`OneFragment`] in a local root, whose next version registers with a mock Tables
/// service.
async fn one_fragment_table() -> (MockServer, Arc<OneFragment>, Table) {
    let server = MockServer::start().await;
    let identity = TableIdentity {
        resource_id: Uuid::new_v4(),
        version_id: Uuid::now_v7(),
    };
    let scandal = Arc::new(OneFragment::new(Uuid::now_v7()));
    Mock::given(method("POST"))
        .and(path(format!(
            "/api/tables/{}/versions",
            identity.resource_id
        )))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "versionId": scandal.next_version,
            "parentVersionId": identity.version_id,
            "seqNo": 2,
        })))
        .mount(&server)
        .await;
    let client = Client::builder()
        .api_url(server.uri())
        .bearer_token("test")
        .scandal(Arc::clone(&scandal) as Arc<dyn Scandal>)
        .build()
        .unwrap();
    let table = client.local_table(&root(), identity).unwrap();
    (server, scandal, table)
}

fn root() -> Url {
    Url::parse("file:///tables/edits/").unwrap()
}

/// Another version of `table`, in the same local root.
fn other_version(table: &Table) -> Table {
    let identity = TableIdentity {
        resource_id: table.identity().resource_id,
        version_id: Uuid::now_v7(),
    };
    let tables = table.connection.tables().unwrap();
    tables.add_local_root(identity, &root()).unwrap();
    Table::new(Arc::clone(&table.connection), identity)
}

/// The kinds, row IDs and labels a change read returns, in order.
async fn read_changes(table: &Table, images: ChangeImages) -> Vec<(String, u64, Option<String>)> {
    let options = ChangeOptions {
        images,
        ..Default::default()
    };
    let scan = table
        .changes_since(&other_version(table), options)
        .await
        .unwrap();
    let batches = scan.into_stream().try_collect::<Vec<_>>().await.unwrap();
    let mut rows = Vec::new();
    for batch in batches {
        let kinds = batch
            .column_by_name(CHANGE_KIND_COLUMN)
            .unwrap()
            .as_string::<i32>();
        let ids = batch
            .column_by_name(ROW_ID_COLUMN)
            .unwrap()
            .as_primitive::<UInt64Type>();
        let labels = batch
            .column_by_name("label")
            .map(|labels| labels.as_string::<i32>());
        for row in 0..batch.num_rows() {
            let label = labels
                .filter(|labels| labels.is_valid(row))
                .map(|labels| labels.value(row).to_owned());
            rows.push((kinds.value(row).to_owned(), ids.value(row), label));
        }
    }
    rows
}

/// A change read reports deletions, insertions and updates from the plan, reading only the
/// images its mode asks for.
#[tokio::test]
async fn a_change_read_composes_the_plan_with_image_reads() {
    let (_server, _scandal, table) = one_fragment_table().await;
    let row =
        |kind: &str, id: u64, label: Option<&str>| (kind.to_owned(), id, label.map(str::to_owned));
    assert_eq!(
        read_changes(&table, ChangeImages::Keys).await,
        [
            row("delete", 10, None),
            row("insert", 12, None),
            row("update_after", 11, None)
        ]
    );
    assert_eq!(
        read_changes(&table, ChangeImages::After).await,
        [
            row("delete", 10, None),
            row("insert", 12, Some("c")),
            row("update_after", 11, Some("b"))
        ]
    );
    assert_eq!(
        read_changes(&table, ChangeImages::BeforeAndAfter).await,
        [
            row("delete", 10, Some("a")),
            row("insert", 12, Some("c")),
            row("update_before", 11, Some("b")),
            row("update_after", 11, Some("b")),
        ]
    );
}

/// A merge sends its patches, deletions and appended rows as parts of one edit, and
/// publishes them as one version.
#[tokio::test]
async fn a_merge_writes_one_edit_of_patches_deletions_and_appended_rows() {
    let (_server, scandal, table) = one_fragment_table().await;
    let merged = table
        .merge_rows(
            rows(vec![1, 2, 5], vec!["a", "B", "e"]),
            &["id"],
            MergeActions {
                matched: Matched::Patch,
                columns: None,
                not_matched: NotMatched::Insert,
                not_matched_by_source: BySource::Delete,
            },
        )
        .await
        .unwrap();
    assert_eq!((merged.inserted, merged.updated, merged.deleted), (1, 1, 1));
    assert_eq!(merged.table.identity().version_id, scandal.next_version);

    let written = scandal.written.lock();
    let [patch] = written.patches.as_slice() else {
        panic!("one patch batch");
    };
    assert_eq!(
        patch
            .column_by_name(ROW_ID_COLUMN)
            .unwrap()
            .as_primitive::<UInt64Type>()
            .values(),
        &[11]
    );
    assert_eq!(
        patch
            .column_by_name("label")
            .unwrap()
            .as_string::<i32>()
            .value(0),
        "B"
    );
    let [deletion] = written.deletions.as_slice() else {
        panic!("one deletion batch");
    };
    assert_eq!(
        deletion.column(0).as_primitive::<UInt64Type>().values(),
        &[12]
    );
    let [append] = written.appends.as_slice() else {
        panic!("one appended batch");
    };
    assert_eq!(append.column(0).as_primitive::<Int64Type>().values(), &[5]);
    let [(base, receipts)] = written.finalized.as_slice() else {
        panic!("one version");
    };
    assert_eq!(
        base,
        &WriteBase::Edit(EditBase {
            version_id: table.identity().version_id,
            schema_changes: None,
        })
    );
    let mut tokens = receipts
        .iter()
        .map(|receipt| receipt.token.as_str())
        .collect::<Vec<_>>();
    tokens.sort_unstable();
    assert_eq!(tokens, ["append", "delete", "patch"]);
}

/// A merge that changes nothing writes nothing and returns the table it started from.
#[tokio::test]
async fn a_merge_without_changes_writes_nothing() {
    let (_server, scandal, table) = one_fragment_table().await;
    let merged = table
        .merge_rows(rows(vec![2], vec!["b"]), &["id"], MergeActions::default())
        .await
        .unwrap();
    assert_eq!(merged.table.identity(), table.identity());
    assert_eq!((merged.inserted, merged.updated, merged.deleted), (0, 0, 0));
    assert!(scandal.written.lock().finalized.is_empty());
}

/// A schema change is an edit with the change and no receipts.
#[tokio::test]
async fn a_schema_change_writes_an_edit_without_receipts() {
    let (_server, scandal, table) = one_fragment_table().await;
    let altered = table
        .alter_schema([SchemaChange::RenameColumn {
            path: ColumnPath::from("label"),
            name: "caption".into(),
        }])
        .await
        .unwrap();
    assert_eq!(altered.identity().version_id, scandal.next_version);
    let written = scandal.written.lock();
    let [(WriteBase::Edit(edit), receipts)] = written.finalized.as_slice() else {
        panic!("one edit");
    };
    assert!(receipts.is_empty());
    assert_eq!(edit.schema_changes.as_ref().map(Vec::len), Some(1));
}