datum-sql 0.10.2

use std::sync::Arc;

use arrow::array::{Array, Int64Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datum::Source;
use datum_sql::{ChangeOp, ChangelogBatch, DatumSqlContext};

#[test]
fn changelog_batch_validates_update_pair_adjacency() {
    let schema = Arc::new(Schema::new(vec![Field::new(
        "account_id",
        DataType::Int64,
        false,
    )]));
    let batch = RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![7, 7]))])
        .expect("input batch builds");

    let changes = ChangelogBatch::try_new(
        vec![ChangeOp::UpdateDelete, ChangeOp::UpdateInsert],
        batch.clone(),
    )
    .expect("adjacent update pair is accepted");
    assert_eq!(
        changes.ops(),
        &[ChangeOp::UpdateDelete, ChangeOp::UpdateInsert]
    );

    let error = ChangelogBatch::try_new(vec![ChangeOp::UpdateInsert, ChangeOp::Insert], batch)
        .expect_err("unpaired update-insert is rejected");
    assert!(
        error
            .to_string()
            .contains("UpdateInsert at row 0 must be immediately preceded")
    );
}

#[tokio::test]
async fn append_only_sink_rejects_updating_stream() {
    let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
    let input = ChangelogBatch::insert_only(
        RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(Int64Array::from(vec![1]))],
        )
        .expect("input batch builds"),
    );

    let context = DatumSqlContext::new();
    context
        .register_changelog_source("updates", schema, Source::from_iter([input]))
        .expect("changelog source registers");
    context
        .register_append_sink("plain_out", |_| Ok(()))
        .expect("append sink registers");

    let error = context
        .select_into("SELECT id FROM updates", "plain_out")
        .await
        .expect_err("updating stream should be rejected by append-only sink");

    assert!(error.to_string().contains(
        "append-only sink cannot consume an updating stream; register a changelog-aware sink"
    ));
}

#[tokio::test]
async fn append_only_sql_execution_stays_plain_record_batches() {
    let schema = Arc::new(Schema::new(vec![
        Field::new("city", DataType::Utf8, false),
        Field::new("temp", DataType::Int64, false),
    ]));
    let batch = RecordBatch::try_new(
        Arc::clone(&schema),
        vec![
            Arc::new(StringArray::from(vec!["sf", "nyc"])),
            Arc::new(Int64Array::from(vec![67, 74])),
        ],
    )
    .expect("input batch builds");

    let context = DatumSqlContext::new();
    context
        .register_source("weather", schema, Source::from_iter([batch]))
        .expect("source registers");

    let output = context
        .execute("SELECT city, temp FROM weather WHERE temp >= 70")
        .await
        .expect("query executes");

    assert_eq!(output.len(), 1);
    assert_eq!(output[0].num_columns(), 2);
    assert_eq!(output[0].num_rows(), 1);
    assert_eq!(string_values(&output[0], 0), vec!["nyc"]);
    assert_eq!(int_values(&output[0], 1), vec![74]);
}

fn int_values(batch: &RecordBatch, column: usize) -> Vec<i64> {
    let array = batch
        .column(column)
        .as_any()
        .downcast_ref::<Int64Array>()
        .expect("column is Int64");
    (0..array.len()).map(|row| array.value(row)).collect()
}

fn string_values(batch: &RecordBatch, column: usize) -> Vec<&str> {
    let array = batch
        .column(column)
        .as_any()
        .downcast_ref::<StringArray>()
        .expect("column is Utf8");
    (0..array.len()).map(|row| array.value(row)).collect()
}