datum-sql 0.10.3

use std::collections::BTreeSet;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

#[cfg(feature = "mq")]
use std::{path::PathBuf, process::Command, sync::mpsc};

use arrow::array::{Array, Int64Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
#[cfg(feature = "mq")]
use datum::{Keep, Sink};
use datum::{NotUsed, Source, StreamError, StreamResult};
#[cfg(feature = "mq")]
use datum_sql::JsonRowFormat;
use datum_sql::{
    ChangeOp, ChangelogBatch, CommittableRecordBatch, ContinuousQueryHandle, DatumSqlContext,
    SourceCommit, SqlSourcePosition,
};

#[tokio::test]
async fn insert_into_append_sink_commits_only_after_sink_confirmation() {
    let schema = test_schema();
    let events = Arc::new(Mutex::new(Vec::<String>::new()));
    let source = Source::from_iter([
        committable_batch(
            Arc::clone(&schema),
            1,
            "alice",
            "batch-1",
            Arc::clone(&events),
        ),
        committable_batch(
            Arc::clone(&schema),
            2,
            "bob",
            "batch-2",
            Arc::clone(&events),
        ),
    ]);

    let context = DatumSqlContext::new();
    context
        .register_committable_source("people", Arc::clone(&schema), source)
        .expect("committable source registers");
    let sink_events = Arc::clone(&events);
    context
        .register_append_sink("out", move |batch| {
            sink_events
                .lock()
                .expect("events lock")
                .push(format!("write:{:?}", int_values(&batch, 0)));
            Ok(())
        })
        .expect("sink registers");

    let mut handle = context
        .execute_insert_into("INSERT INTO out SELECT id, name FROM people WHERE id >= 2")
        .await
        .expect("insert materializes");
    wait_for_completion(&mut handle, Duration::from_secs(2)).expect("insert completes");

    assert_eq!(
        events.lock().expect("events lock").as_slice(),
        ["write:[]", "commit:batch-1", "write:[2]", "commit:batch-2"]
    );
}

#[tokio::test]
async fn sink_failure_leaves_source_position_uncommitted() {
    let schema = test_schema();
    let events = Arc::new(Mutex::new(Vec::<String>::new()));
    let source = Source::from_iter([committable_batch(
        Arc::clone(&schema),
        7,
        "carol",
        "batch-7",
        Arc::clone(&events),
    )]);

    let context = DatumSqlContext::new();
    context
        .register_committable_source("people", schema, source)
        .expect("source registers");
    let sink_events = Arc::clone(&events);
    context
        .register_append_sink("out", move |batch| {
            sink_events
                .lock()
                .expect("events lock")
                .push(format!("write:{:?}", int_values(&batch, 0)));
            Err(StreamError::Failed("sink rejected batch".to_owned()))
        })
        .expect("sink registers");

    let mut handle = context
        .execute_insert_into("INSERT INTO out SELECT id, name FROM people")
        .await
        .expect("insert materializes");
    let result = wait_for_completion(&mut handle, Duration::from_secs(2));

    assert!(result.is_err(), "sink failure should fail the query");
    assert_eq!(
        events.lock().expect("events lock").as_slice(),
        ["write:[7]"]
    );
}

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

    let context = DatumSqlContext::new();
    context
        .register_changelog_source("changes", schema, Source::from_iter([changes]))
        .expect("changelog source registers");
    context
        .register_append_sink("out", |_batch| Ok(()))
        .expect("append sink registers");

    let error = context
        .execute_insert_into("INSERT INTO out SELECT id FROM changes")
        .await
        .expect_err("append sink rejects updating stream");

    assert!(
        error
            .to_string()
            .contains("cannot consume an updating stream")
    );
}

#[tokio::test]
async fn insert_into_allows_changelog_aware_sink() {
    let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
    let changes = ChangelogBatch::try_new(
        vec![ChangeOp::Insert, ChangeOp::Delete],
        RecordBatch::try_new(
            Arc::clone(&schema),
            vec![Arc::new(Int64Array::from(vec![1, 1]))],
        )
        .expect("batch builds"),
    )
    .expect("changelog builds");
    let seen_ops = Arc::new(Mutex::new(Vec::<Vec<ChangeOp>>::new()));

    let context = DatumSqlContext::new();
    context
        .register_changelog_source("changes", schema, Source::from_iter([changes]))
        .expect("changelog source registers");
    let seen_ops_for_sink = Arc::clone(&seen_ops);
    context
        .register_changelog_sink(
            "out",
            |_batch| Ok(()),
            move |changes| {
                seen_ops_for_sink
                    .lock()
                    .expect("ops lock")
                    .push(changes.ops().to_vec());
                Ok(())
            },
        )
        .expect("changelog sink registers");

    let mut handle = context
        .execute_insert_into("INSERT INTO out SELECT id FROM changes")
        .await
        .expect("insert materializes");
    wait_for_completion(&mut handle, Duration::from_secs(2)).expect("insert completes");

    assert_eq!(
        seen_ops.lock().expect("ops lock").as_slice(),
        [vec![ChangeOp::Insert, ChangeOp::Delete]]
    );
}

#[cfg(feature = "mq")]
#[tokio::test]
async fn kafka_insert_into_replays_after_forced_restart_when_env_is_set() {
    if std::env::var_os("DATUM_SQL_CONNECTOR_INTEGRATION").is_none() {
        return;
    }

    let Some(bootstrap) = std::env::var("MQ_BOOTSTRAP_SERVERS").ok() else {
        eprintln!("skipping Kafka sink integration: MQ_BOOTSTRAP_SERVERS is not set");
        return;
    };

    let source_topic = format!("datum-sql-sink-src-{}", unique_suffix());
    let sink_topic = format!("datum-sql-sink-out-{}", unique_suffix());
    create_kafka_topic(&bootstrap, &source_topic);
    create_kafka_topic(&bootstrap, &sink_topic);

    let schema = test_schema();
    let format = JsonRowFormat::new(Arc::clone(&schema));
    produce_json_rows(
        &bootstrap,
        &source_topic,
        [
            br#"{"id":1,"name":"one"}"#.as_slice(),
            br#"{"id":2,"name":"two"}"#.as_slice(),
            br#"{"id":3,"name":"three"}"#.as_slice(),
        ],
    );

    let group = format!("datum-sql-sink-group-{}", unique_suffix());
    let source_settings = kafka_source_settings(&bootstrap, &group);
    let producer_settings = datum_mq::KafkaProducerSettings::new(bootstrap.clone());

    let first = DatumSqlContext::new();
    first
        .register_mq_topic(
            "people",
            source_settings.clone(),
            source_topic.clone(),
            format.clone(),
        )
        .expect("source registers");
    let failing_topic = sink_topic.clone();
    let failing_bootstrap = bootstrap.clone();
    let failing_format = format.clone();
    first
        .register_append_sink("out", move |batch| {
            write_json_batch_to_kafka(&failing_bootstrap, &failing_topic, &failing_format, &batch)?;
            Err(StreamError::Failed(
                "forced restart after sink confirmation before source commit".to_owned(),
            ))
        })
        .expect("failing sink registers");

    let mut first_handle = first
        .execute_insert_into("INSERT INTO out SELECT id, name FROM people")
        .await
        .expect("first insert materializes");
    let first_result = wait_for_completion(&mut first_handle, Duration::from_secs(30));
    assert!(
        first_result.is_err(),
        "forced restart should fail the first query before commit"
    );

    let second = DatumSqlContext::new();
    second
        .register_mq_topic("people", source_settings, source_topic, format.clone())
        .expect("source registers");
    second
        .register_mq_json_sink("out", producer_settings, sink_topic.clone(), format.clone())
        .expect("Kafka sink registers");

    let second_handle = second
        .execute_insert_into("INSERT INTO out SELECT id, name FROM people")
        .await
        .expect("second insert materializes");

    let rows = read_kafka_rows_until(
        &bootstrap,
        &sink_topic,
        Arc::clone(&schema),
        4,
        Duration::from_secs(30),
    );
    second_handle.cancel();
    drop(second_handle);

    let ids = rows.iter().map(|row| row.0).collect::<Vec<_>>();
    let unique = ids.iter().copied().collect::<BTreeSet<_>>();
    assert_eq!(unique, BTreeSet::from([1, 2, 3]));
    assert!(
        rows.len() >= 4,
        "restart should replay at least one already-produced row, got {rows:?}"
    );
}

fn committable_batch(
    schema: SchemaRef,
    id: i64,
    name: &str,
    label: &'static str,
    events: Arc<Mutex<Vec<String>>>,
) -> CommittableRecordBatch {
    let batch = RecordBatch::try_new(
        schema,
        vec![
            Arc::new(Int64Array::from(vec![id])),
            Arc::new(StringArray::from(vec![name])),
        ],
    )
    .expect("batch builds");
    let commit = SourceCommit::from_fn(label, move || {
        events
            .lock()
            .expect("events lock")
            .push(format!("commit:{label}"));
        Ok(())
    });
    CommittableRecordBatch::new(
        batch,
        Some(SqlSourcePosition::custom("test", label)),
        0,
        commit,
    )
}

fn test_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("name", DataType::Utf8, false),
    ]))
}

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()
}

#[cfg(feature = "mq")]
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()
}

fn wait_for_completion(
    handle: &mut ContinuousQueryHandle,
    timeout: Duration,
) -> StreamResult<NotUsed> {
    let deadline = Instant::now() + timeout;
    loop {
        if let Some(result) = handle.try_wait() {
            return result;
        }
        assert!(
            Instant::now() < deadline,
            "query did not complete within {timeout:?}"
        );
        thread::yield_now();
    }
}

#[cfg(feature = "mq")]
fn kafka_source_settings(bootstrap: &str, group: &str) -> datum_mq::KafkaConsumerSettings {
    let mut settings = datum_mq::KafkaConsumerSettings::new(bootstrap, group)
        .with("auto.offset.reset", "earliest")
        .with_backpressure(8, 16)
        .with_poll_batch_size(16)
        .with_commit_batch_size(1)
        .with_commit_interval(Duration::from_millis(10));
    settings.commit_sync = true;
    settings
}

#[cfg(feature = "mq")]
fn produce_json_rows<'a, I>(bootstrap: &str, topic: &str, rows: I)
where
    I: IntoIterator<Item = &'a [u8]>,
{
    let records = rows.into_iter().map(|payload| {
        datum_mq::ProducerRecord::new(topic.to_owned(), bytes::Bytes::copy_from_slice(payload))
    });
    let producer = Source::from_iter(records)
        .run_with(datum_mq::KafkaSink::plain(
            datum_mq::KafkaProducerSettings::new(bootstrap.to_owned()),
        ))
        .expect("Kafka producer materializes");
    producer
        .drain_and_shutdown()
        .expect("Kafka producer drains");
}

#[cfg(feature = "mq")]
fn write_json_batch_to_kafka(
    bootstrap: &str,
    topic: &str,
    format: &JsonRowFormat,
    batch: &RecordBatch,
) -> StreamResult<()> {
    let payloads = format.encode_record_batch(batch).map_err(to_stream_error)?;
    let records = payloads.into_iter().map(|payload| {
        datum_mq::ProducerRecord::new(topic.to_owned(), bytes::Bytes::from(payload))
    });
    let producer = Source::from_iter(records)
        .run_with(datum_mq::KafkaSink::plain(
            datum_mq::KafkaProducerSettings::new(bootstrap.to_owned()),
        ))
        .map_err(to_stream_error)?;
    producer.drain_and_shutdown().map_err(to_stream_error)
}

#[cfg(feature = "mq")]
fn read_kafka_rows_until(
    bootstrap: &str,
    topic: &str,
    schema: SchemaRef,
    min_rows: usize,
    timeout: Duration,
) -> Vec<(i64, String)> {
    let group = format!("datum-sql-sink-read-{}", unique_suffix());
    let settings = datum_mq::KafkaConsumerSettings::new(bootstrap.to_owned(), group)
        .with("auto.offset.reset", "earliest")
        .with_backpressure(8, 16)
        .with_poll_batch_size(16);
    let (tx, rx) = mpsc::channel();
    let (control, completion) = datum_sql::connect::mq::kafka_json_source(
        settings,
        topic.to_owned(),
        JsonRowFormat::new(schema),
    )
    .to_mat(
        Sink::foreach(move |batch: datum_sql::KafkaRecordBatch| {
            for row in int_values(batch.batch(), 0).into_iter().zip(
                string_values(batch.batch(), 1)
                    .into_iter()
                    .map(str::to_owned),
            ) {
                let _ = tx.send(row);
            }
            let _ = batch.batch().num_rows();
        }),
        Keep::both,
    )
    .run()
    .expect("Kafka reader materializes");

    let mut rows = Vec::new();
    let deadline = Instant::now() + timeout;
    while rows.len() < min_rows && Instant::now() < deadline {
        if let Ok(row) = rx.recv_timeout(Duration::from_millis(250)) {
            rows.push(row);
        }
    }

    control.shutdown_now();
    let _ = completion.wait();
    rows
}

#[cfg(feature = "mq")]
fn to_stream_error(error: impl std::fmt::Display) -> StreamError {
    StreamError::Failed(error.to_string())
}

#[cfg(feature = "mq")]
fn create_kafka_topic(bootstrap: &str, topic: &str) {
    let topics = kafka_topics_bin();
    let status = Command::new(&topics)
        .args([
            "--bootstrap-server",
            bootstrap,
            "--create",
            "--if-not-exists",
            "--topic",
            topic,
            "--partitions",
            "1",
            "--replication-factor",
            "1",
        ])
        .status()
        .unwrap_or_else(|error| panic!("failed to run {}: {error}", topics.display()));
    assert!(status.success(), "Kafka topic creation failed: {status}");
}

#[cfg(feature = "mq")]
fn kafka_topics_bin() -> PathBuf {
    if let Some(path) = std::env::var_os("KAFKA_TOPICS_BIN") {
        return PathBuf::from(path);
    }
    repo_root().join("baselines/mq/.runtime/kafka_2.13-4.2.0/bin/kafka-topics.sh")
}

#[cfg(feature = "mq")]
fn repo_root() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..")
}

#[cfg(feature = "mq")]
fn unique_suffix() -> u128 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .expect("system time is after unix epoch")
        .as_nanos()
        ^ u128::from(std::process::id())
}