datum-sql 0.10.3

use std::sync::Arc;

#[cfg(feature = "cdc")]
use std::sync::Mutex;
#[cfg(any(feature = "mq", feature = "cdc"))]
use std::sync::mpsc;
#[cfg(any(feature = "mq", feature = "cdc"))]
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
#[cfg(feature = "mq")]
use std::{path::PathBuf, process::Command};

#[cfg(any(feature = "mq", feature = "cdc"))]
use arrow::array::Array;
use arrow::array::{Int64Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use datum::Source;
#[cfg(any(feature = "mq", feature = "cdc"))]
use datum::{Keep, Sink};
use datum_sql::{ChangelogBatch, DatumSqlContext};

#[cfg(feature = "mq")]
use datum_sql::{
    BatchEnvelope, EnvelopedRecordBatch, JsonRowFormat, KafkaPartitionOffset, KafkaSourcePosition,
    MqPayloadFormat,
};

#[cfg(feature = "cdc")]
use datum_cdc::{
    CdcOffset, CdcSource, ChangeEvent, ChangeOperation, ColumnMetadata, ColumnValue, PgLsn,
    RelationMetadata, ReplicaIdentity, RowData, SlotLifecycle, SourceMetadata, TransactionMeta,
};
#[cfg(feature = "cdc")]
use datum_sql::{CdcChangelogBatch, CdcTableAdapter, ChangeOp};

#[cfg(feature = "mq")]
#[test]
fn mq_json_format_decodes_rows_and_keeps_position_envelope() {
    let format = JsonRowFormat::new(test_schema()).with_schema_revision(9);
    let batch = format
        .decode_payload_slices([
            br#"{"id":1,"name":"alice"}"#.as_slice(),
            br#"{"id":2,"name":"bob"}"#.as_slice(),
        ])
        .expect("JSON payloads decode");

    assert_eq!(int_values(&batch, 0), vec![1, 2]);
    assert_eq!(string_values(&batch, 1), vec!["alice", "bob"]);
    assert_eq!(format.schema_revision(), 9);

    let position = KafkaSourcePosition::from_offsets([KafkaPartitionOffset {
        topic: "bids".to_owned(),
        partition: 3,
        first_offset: 41,
        last_offset: 43,
    }]);
    let enveloped = EnvelopedRecordBatch::new(
        batch,
        BatchEnvelope::new(position, format.schema_revision()),
    );

    assert_eq!(enveloped.envelope().schema_revision(), 9);
    assert_eq!(
        enveloped.envelope().source_position().offsets()[0].next_offset(),
        44
    );
}

#[cfg(feature = "mq")]
#[tokio::test]
async fn register_mq_topic_plans_without_materializing_kafka() {
    let context = DatumSqlContext::new();
    let settings = datum_mq::KafkaConsumerSettings::new("127.0.0.1:9092", "datum-sql-test");
    context
        .register_mq_topic(
            "bids",
            settings,
            "datum-sql-test-bids",
            JsonRowFormat::new(test_schema()),
        )
        .expect("MQ table registers");

    context
        .physical_plan("SELECT id, name FROM bids")
        .await
        .expect("registered MQ topic plans");
}

#[cfg(feature = "cdc")]
#[test]
fn cdc_adapter_maps_insert_delete_and_update_pairs() {
    let adapter = CdcTableAdapter::new(test_schema());

    let insert = adapter
        .map_change_event(change_event(
            ChangeOperation::Insert,
            None,
            Some(row(1, "alice")),
            3,
            0,
        ))
        .expect("insert maps");
    assert_eq!(insert.batch().ops(), &[ChangeOp::Insert]);
    assert_eq!(insert.envelope().schema_revision(), 3);
    assert_eq!(insert.envelope().source_position().lsn().event_index, 0);
    assert_eq!(int_values(insert.batch().batch(), 0), vec![1]);

    let delete = adapter
        .map_change_event(change_event(
            ChangeOperation::Delete,
            Some(row(1, "alice")),
            None,
            3,
            1,
        ))
        .expect("delete maps");
    assert_eq!(delete.batch().ops(), &[ChangeOp::Delete]);
    assert_eq!(string_values(delete.batch().batch(), 1), vec!["alice"]);

    let update = adapter
        .map_change_event(change_event(
            ChangeOperation::Update,
            Some(row(1, "alice")),
            Some(row(1, "alicia")),
            4,
            2,
        ))
        .expect("update maps");
    assert_eq!(
        update.batch().ops(),
        &[ChangeOp::UpdateDelete, ChangeOp::UpdateInsert]
    );
    assert_eq!(update.envelope().schema_revision(), 4);
    assert_eq!(
        string_values(update.batch().batch(), 1),
        vec!["alice", "alicia"]
    );
}

#[cfg(feature = "cdc")]
#[test]
fn cdc_adapter_uses_unpaired_delete_insert_without_before_image() {
    let adapter = CdcTableAdapter::new(test_schema());
    let update = adapter
        .map_change_event(change_event(
            ChangeOperation::Update,
            None,
            Some(row(7, "visible")),
            5,
            0,
        ))
        .expect("N-only update maps");

    assert_eq!(update.batch().ops(), &[ChangeOp::Delete, ChangeOp::Insert]);
    assert_eq!(
        string_values(update.batch().batch(), 1),
        vec!["visible", "visible"]
    );
}

#[tokio::test]
async fn changelog_catalog_rejects_append_only_execution() {
    let batch = ChangelogBatch::insert_only(
        RecordBatch::try_new(
            test_schema(),
            vec![
                Arc::new(Int64Array::from(vec![1])),
                Arc::new(StringArray::from(vec!["alice"])),
            ],
        )
        .expect("batch builds"),
    );

    let context = DatumSqlContext::new();
    context
        .register_changelog_source("accounts", test_schema(), Source::from_iter([batch]))
        .expect("changelog source registers");

    let error = context
        .execute("SELECT id FROM accounts")
        .await
        .expect_err("append-only execution rejects updating table");
    assert!(error.to_string().contains("updating tables produce"));

    context
        .register_append_sink("plain_out", |_| Ok(()))
        .expect("append sink registers");
    let error = context
        .select_into("SELECT id FROM accounts", "plain_out")
        .await
        .expect_err("append-only sink rejects updating table");
    assert!(error.to_string().contains(
        "append-only sink cannot consume an updating stream; register a changelog-aware sink"
    ));
}

#[cfg(feature = "mq")]
#[test]
fn mq_adapter_consumes_real_kafka_when_integration_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 integration: MQ_BOOTSTRAP_SERVERS is not set");
        return;
    };

    let topic = format!("datum-sql-it-{}", unique_suffix());
    create_kafka_topic(&bootstrap, &topic);

    let records = vec![
        datum_mq::ProducerRecord::new(
            topic.clone(),
            Some(bytes::Bytes::from_static(br#"{"id":101,"name":"kafka-a"}"#)),
        ),
        datum_mq::ProducerRecord::new(
            topic.clone(),
            Some(bytes::Bytes::from_static(br#"{"id":102,"name":"kafka-b"}"#)),
        ),
    ];
    let producer = Source::from_iter(records)
        .run_with(datum_mq::KafkaSink::plain(
            datum_mq::KafkaProducerSettings::new(bootstrap.clone()),
        ))
        .expect("Kafka producer materializes");
    producer
        .drain_and_shutdown()
        .expect("Kafka producer drains");

    let group = format!("datum-sql-it-{}", unique_suffix());
    let settings = datum_mq::KafkaConsumerSettings::new(bootstrap, 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,
        JsonRowFormat::new(test_schema()).with_schema_revision(17),
    )
    .to_mat(
        Sink::foreach(move |batch| {
            let _ = tx.send(batch);
        }),
        Keep::both,
    )
    .run()
    .expect("Kafka SQL adapter materializes");

    let mut rows = Vec::new();
    let mut schema_revision = None;
    let mut saw_position = false;
    let deadline = Instant::now() + Duration::from_secs(30);
    while rows.len() < 2 && Instant::now() < deadline {
        if let Ok(batch) = rx.recv_timeout(Duration::from_millis(250)) {
            schema_revision = Some(batch.envelope().schema_revision());
            saw_position |= !batch.envelope().source_position().is_empty();
            rows.extend(
                int_values(batch.batch(), 0).into_iter().zip(
                    string_values(batch.batch(), 1)
                        .into_iter()
                        .map(str::to_owned),
                ),
            );
        }
    }

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

    rows.sort_by_key(|row| row.0);
    assert_eq!(
        rows,
        vec![(101, "kafka-a".to_owned()), (102, "kafka-b".to_owned())]
    );
    assert_eq!(schema_revision, Some(17));
    assert!(saw_position, "Kafka adapter should expose source offsets");
}

#[cfg(feature = "mq-native")]
#[test]
fn mq_native_backend_sql_kafka_smoke_matches_rdkafka_checksum() {
    if std::env::var_os("DATUM_SQL_NATIVE_KAFKA_SMOKE").is_none() {
        return;
    }

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

    let topic = format!("datum-sql-native-q0-{}", unique_suffix());
    create_kafka_topic(&bootstrap, &topic);
    let records = vec![
        datum_mq::ProducerRecord::new(
            topic.clone(),
            Some(bytes::Bytes::from_static(br#"{"id":201,"name":"q0-a"}"#)),
        ),
        datum_mq::ProducerRecord::new(
            topic.clone(),
            Some(bytes::Bytes::from_static(br#"{"id":202,"name":"q0-b"}"#)),
        ),
        datum_mq::ProducerRecord::new(
            topic.clone(),
            Some(bytes::Bytes::from_static(br#"{"id":203,"name":"q0-c"}"#)),
        ),
    ];
    let producer = Source::from_iter(records)
        .run_with(datum_mq::KafkaSink::plain(
            datum_mq::KafkaProducerSettings::new(bootstrap.clone()),
        ))
        .expect("Kafka producer materializes");
    producer
        .drain_and_shutdown()
        .expect("Kafka producer drains");

    let stock = read_kafka_sql_rows(
        &bootstrap,
        &topic,
        "stock",
        datum_mq::KafkaConsumerBackend::Rdkafka,
        3,
    );
    let native = read_kafka_sql_rows(
        &bootstrap,
        &topic,
        "native",
        datum_mq::KafkaConsumerBackend::Native,
        3,
    );
    assert_eq!(native, stock);
    assert_eq!(rows_checksum(&native), rows_checksum(&stock));
}

#[cfg(feature = "mq-native")]
fn read_kafka_sql_rows(
    bootstrap: &str,
    topic: &str,
    suffix: &str,
    backend: datum_mq::KafkaConsumerBackend,
    expected: usize,
) -> Vec<(i64, String)> {
    let group = format!("datum-sql-native-q0-{suffix}-{}", unique_suffix());
    let settings = datum_mq::KafkaConsumerSettings::new(bootstrap.to_owned(), group)
        .with_consumer_backend(backend)
        .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(test_schema()).with_schema_revision(18),
    )
    .to_mat(
        Sink::foreach(move |batch| {
            let _ = tx.send(batch);
        }),
        Keep::both,
    )
    .run()
    .expect("Kafka SQL adapter materializes");

    let mut rows = Vec::new();
    let deadline = Instant::now() + Duration::from_secs(30);
    while rows.len() < expected && Instant::now() < deadline {
        if let Ok(batch) = rx.recv_timeout(Duration::from_millis(250)) {
            rows.extend(
                int_values(batch.batch(), 0).into_iter().zip(
                    string_values(batch.batch(), 1)
                        .into_iter()
                        .map(str::to_owned),
                ),
            );
        }
    }
    control.shutdown_now();
    let _ = completion.wait();
    rows.sort_by_key(|row| row.0);
    rows
}

#[cfg(feature = "mq-native")]
fn rows_checksum(rows: &[(i64, String)]) -> u64 {
    let mut hash = 0xcbf29ce484222325_u64;
    for (id, name) in rows {
        for byte in id
            .to_le_bytes()
            .into_iter()
            .chain(name.as_bytes().iter().copied())
        {
            hash ^= u64::from(byte);
            hash = hash.wrapping_mul(0x100000001b3);
        }
    }
    hash
}

#[cfg(feature = "cdc")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn cdc_adapter_consumes_real_postgres_when_integration_env_is_set() {
    if std::env::var_os("DATUM_SQL_CONNECTOR_INTEGRATION").is_none() {
        return;
    }

    let Some(database_url) = std::env::var("CDC_POSTGRES_URL").ok() else {
        eprintln!("skipping CDC integration: CDC_POSTGRES_URL is not set");
        return;
    };

    let publication =
        std::env::var("CDC_PUBLICATION").unwrap_or_else(|_| "datum_cdc_pub".to_owned());
    let slot = format!("datum_sql_it_{}", unique_suffix());
    let run_id = format!("datum-sql-it-{}", unique_suffix());
    let row_id = (unique_suffix() % 1_000_000_000) as i64;

    let source = CdcSource::postgres()
        .connect_url(&database_url)
        .expect("CDC URL parses")
        .slot(slot)
        .publication(publication)
        .slot_lifecycle(SlotLifecycle::CreateOwned)
        .status_interval(Duration::from_millis(50))
        .idle_wakeup_interval(Duration::from_millis(50))
        .build()
        .expect("CDC source builds");

    let (tx, rx) = mpsc::channel();
    let checkpoint = Arc::new(Mutex::new(None::<datum_cdc::CdcCheckpointHandle>));
    let checkpoint_for_sink = Arc::clone(&checkpoint);
    let run_id_for_sink = run_id.clone();
    let (handle, completion) =
        datum_sql::connect::cdc::cdc_changelog_source(source, cdc_events_schema())
            .to_mat(
                Sink::foreach(move |batch: CdcChangelogBatch| {
                    if batch_contains_run(&batch, &run_id_for_sink) {
                        let _ = tx.send(batch.clone());
                    }
                    if let Some(checkpoint) = checkpoint_for_sink
                        .lock()
                        .expect("checkpoint lock poisoned")
                        .clone()
                    {
                        let _ =
                            checkpoint.checkpoint(batch.envelope().source_position().lsn().clone());
                    }
                }),
                Keep::both,
            )
            .run()
            .expect("CDC SQL adapter materializes");
    *checkpoint.lock().expect("checkpoint lock poisoned") = Some(handle.checkpoint_handle());

    tokio::time::sleep(Duration::from_millis(500)).await;
    let (client, connection) = tokio_postgres::connect(&database_url, tokio_postgres::NoTls)
        .await
        .expect("PostgreSQL connects");
    tokio::spawn(async move {
        let _ = connection.await;
    });
    client
        .execute(
            "INSERT INTO public.cdc_events (id, run_id, value, op_seq, kind, commit_ns) \
             VALUES ($1, $2, $3, $4, $5, $6)",
            &[&row_id, &run_id, &11_i64, &1_i64, &"datum-sql", &0_i64],
        )
        .await
        .expect("insert publishes");
    client
        .execute(
            "UPDATE public.cdc_events SET value = $1, op_seq = $2 WHERE id = $3",
            &[&12_i64, &2_i64, &row_id],
        )
        .await
        .expect("update publishes");
    client
        .execute("DELETE FROM public.cdc_events WHERE id = $1", &[&row_id])
        .await
        .expect("delete publishes");

    let mut batches = Vec::new();
    let deadline = Instant::now() + Duration::from_secs(30);
    while batches.len() < 3 && Instant::now() < deadline {
        if let Ok(batch) = rx.recv_timeout(Duration::from_millis(250)) {
            batches.push(batch);
        }
    }

    handle.force_drop_slot().await.expect("CDC slot drops");
    let _ = completion.wait();

    let ops = batches
        .iter()
        .map(|batch| batch.batch().ops().to_vec())
        .collect::<Vec<_>>();
    assert!(
        ops.iter().any(|ops| ops == &[ChangeOp::Insert]),
        "missing insert batch: {ops:?}"
    );
    assert!(
        ops.iter()
            .any(|ops| ops == &[ChangeOp::UpdateDelete, ChangeOp::UpdateInsert]),
        "missing update pair batch: {ops:?}"
    );
    assert!(
        ops.iter().any(|ops| ops == &[ChangeOp::Delete]),
        "missing delete batch: {ops:?}"
    );
    assert!(
        batches
            .iter()
            .all(|batch| batch.envelope().schema_revision() > 0),
        "CDC adapter should carry relation revisions"
    );
}

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

#[cfg(feature = "cdc")]
fn cdc_events_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("run_id", DataType::Utf8, false),
        Field::new("value", DataType::Int64, false),
    ]))
}

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

#[cfg(feature = "cdc")]
fn batch_contains_run(batch: &CdcChangelogBatch, run_id: &str) -> bool {
    string_values(batch.batch().batch(), 1)
        .into_iter()
        .any(|value| value == run_id)
}

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

#[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 = "cdc")]
fn row(id: i64, name: &str) -> RowData {
    RowData {
        values: vec![
            ColumnValue::Text(id.to_string()),
            ColumnValue::Text(name.to_owned()),
        ],
    }
}

#[cfg(feature = "cdc")]
fn change_event(
    op: ChangeOperation,
    before: Option<RowData>,
    after: Option<RowData>,
    revision: u64,
    event_index: u32,
) -> ChangeEvent {
    ChangeEvent {
        source: SourceMetadata {
            database: "db".to_owned(),
            slot: "slot".to_owned(),
            publication: "pub".to_owned(),
        },
        schema: "public".to_owned(),
        table: "accounts".to_owned(),
        op,
        before,
        after,
        truncate: None,
        lsn: CdcOffset {
            slot: "slot".to_owned(),
            tx_end_lsn: PgLsn::from_u64(20),
            commit_lsn: PgLsn::from_u64(10),
            xid: 99,
            event_index,
            event_count: 3,
        },
        tx: TransactionMeta {
            xid: 99,
            commit_time_micros: 123,
            event_index,
            event_count: 3,
        },
        relation: RelationMetadata {
            oid: 42,
            schema: "public".to_owned(),
            table: "accounts".to_owned(),
            replica_identity: ReplicaIdentity::Full,
            columns: vec![
                ColumnMetadata {
                    name: "id".to_owned(),
                    type_oid: 20,
                    type_modifier: -1,
                    key: true,
                },
                ColumnMetadata {
                    name: "name".to_owned(),
                    type_oid: 25,
                    type_modifier: -1,
                    key: false,
                },
            ],
            revision,
        },
    }
}