datum-sql 0.10.3

use std::sync::Arc;

use arrow::array::{
    ArrayRef, BooleanBuilder, Float32Builder, Float64Builder, Int8Builder, Int16Builder,
    Int32Builder, Int64Builder, LargeStringBuilder, StringBuilder, UInt8Builder, UInt16Builder,
    UInt32Builder, UInt64Builder,
};
use arrow::datatypes::{DataType, Field, SchemaRef};
use arrow::record_batch::RecordBatch;
use datafusion::common::{DataFusionError, Result};
use datum::{NotUsed, Source};
use datum_cdc::{
    CdcHandle, CdcOffset, ChangeEvent, ChangeOperation, ColumnValue, RelationMetadata, RowData,
};

use crate::{ChangeOp, ChangelogBatch};

use super::{BatchEnvelope, EnvelopedBatch};

/// Source position for one CDC-backed SQL changelog batch.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CdcSourcePosition {
    lsn: CdcOffset,
}

impl CdcSourcePosition {
    /// Creates a CDC source position from an LSN offset.
    #[must_use]
    pub const fn new(lsn: CdcOffset) -> Self {
        Self { lsn }
    }

    /// Returns the CDC LSN offset.
    #[must_use]
    pub const fn lsn(&self) -> &CdcOffset {
        &self.lsn
    }

    /// Consumes this position and returns the CDC LSN offset.
    #[must_use]
    pub fn into_lsn(self) -> CdcOffset {
        self.lsn
    }
}

/// CDC-backed changelog batch with source position metadata.
pub type CdcChangelogBatch = EnvelopedBatch<ChangelogBatch, CdcSourcePosition>;

/// Maps `datum-cdc` pgoutput change events into typed SQL changelog batches.
#[derive(Debug, Clone)]
pub struct CdcTableAdapter {
    schema: SchemaRef,
}

impl CdcTableAdapter {
    /// Creates a CDC table adapter for a user-supplied Arrow schema.
    #[must_use]
    pub fn new(schema: SchemaRef) -> Self {
        Self { schema }
    }

    /// Returns the Arrow schema produced by this adapter.
    #[must_use]
    pub fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    /// Converts one `datum-cdc` change event into a SQL changelog batch.
    pub fn map_change_event(&self, event: ChangeEvent) -> Result<CdcChangelogBatch> {
        let schema_revision = event.relation.revision;
        let position = CdcSourcePosition::new(event.lsn.clone());
        let (ops, rows) = match event.op {
            ChangeOperation::Insert => {
                let after = required_row(event.after.as_ref(), "insert after row")?;
                (vec![ChangeOp::Insert], vec![after])
            }
            ChangeOperation::Delete => {
                let before = required_row(event.before.as_ref(), "delete before row")?;
                (vec![ChangeOp::Delete], vec![before])
            }
            ChangeOperation::Update => {
                let after = required_row(event.after.as_ref(), "update after row")?;
                if let Some(before) = event.before.as_ref() {
                    (
                        vec![ChangeOp::UpdateDelete, ChangeOp::UpdateInsert],
                        vec![before, after],
                    )
                } else {
                    // A pgoutput N-only update has no old row. Keep the update
                    // unpaired so downstream stages do not infer adjacency.
                    (vec![ChangeOp::Delete, ChangeOp::Insert], vec![after, after])
                }
            }
            ChangeOperation::Truncate => {
                return Err(DataFusionError::Plan(
                    "CDC truncate events cannot be represented as row-level ChangelogBatch values"
                        .into(),
                ));
            }
        };

        let batch = rows_to_record_batch(Arc::clone(&self.schema), &event.relation, &rows)?;
        let changes = ChangelogBatch::try_new(ops, batch)?;
        Ok(CdcChangelogBatch::new(
            changes,
            BatchEnvelope::new(position, schema_revision),
        ))
    }
}

#[must_use]
/// Maps a `datum-cdc` source into SQL changelog batches with source metadata.
pub fn cdc_changelog_source(
    source: Source<ChangeEvent, CdcHandle>,
    schema: SchemaRef,
) -> Source<CdcChangelogBatch, CdcHandle> {
    let adapter = CdcTableAdapter::new(schema);
    source.try_map(move |event| adapter.map_change_event(event).map_err(crate::stream_error))
}

#[must_use]
/// Maps a `datum-cdc` source into plain SQL changelog payloads.
pub fn cdc_changelog_batch_source_uncontrolled(
    source: Source<ChangeEvent, CdcHandle>,
    schema: SchemaRef,
) -> Source<ChangelogBatch, NotUsed> {
    cdc_changelog_source(source, schema)
        .map(|batch| batch.into_batch())
        .map_materialized_value(|_| NotUsed)
}

fn required_row<'a>(row: Option<&'a RowData>, context: &str) -> Result<&'a RowData> {
    row.ok_or_else(|| DataFusionError::Plan(format!("CDC event is missing {context}")))
}

fn rows_to_record_batch(
    schema: SchemaRef,
    relation: &RelationMetadata,
    rows: &[&RowData],
) -> Result<RecordBatch> {
    if rows.is_empty() {
        return Ok(RecordBatch::new_empty(schema));
    }

    let mut builders = schema
        .fields()
        .iter()
        .map(|field| CdcColumnBuilder::try_new(field, rows.len()))
        .collect::<Result<Vec<_>>>()?;

    for row in rows {
        for (field_index, field) in schema.fields().iter().enumerate() {
            let value = row.get(relation, field.name()).ok_or_else(|| {
                DataFusionError::Plan(format!(
                    "CDC row for {}.{} does not contain column {:?}",
                    relation.schema,
                    relation.table,
                    field.name()
                ))
            })?;
            builders[field_index].append(field, value)?;
        }
    }

    let columns = builders
        .into_iter()
        .map(CdcColumnBuilder::finish)
        .collect::<Vec<_>>();
    RecordBatch::try_new(schema, columns).map_err(DataFusionError::from)
}

enum CdcColumnBuilder {
    Boolean(BooleanBuilder),
    Int8(Int8Builder),
    Int16(Int16Builder),
    Int32(Int32Builder),
    Int64(Int64Builder),
    UInt8(UInt8Builder),
    UInt16(UInt16Builder),
    UInt32(UInt32Builder),
    UInt64(UInt64Builder),
    Float32(Float32Builder),
    Float64(Float64Builder),
    Utf8(StringBuilder),
    LargeUtf8(LargeStringBuilder),
}

impl CdcColumnBuilder {
    fn try_new(field: &Field, rows: usize) -> Result<Self> {
        match field.data_type() {
            DataType::Boolean => Ok(Self::Boolean(BooleanBuilder::with_capacity(rows))),
            DataType::Int8 => Ok(Self::Int8(Int8Builder::with_capacity(rows))),
            DataType::Int16 => Ok(Self::Int16(Int16Builder::with_capacity(rows))),
            DataType::Int32 => Ok(Self::Int32(Int32Builder::with_capacity(rows))),
            DataType::Int64 => Ok(Self::Int64(Int64Builder::with_capacity(rows))),
            DataType::UInt8 => Ok(Self::UInt8(UInt8Builder::with_capacity(rows))),
            DataType::UInt16 => Ok(Self::UInt16(UInt16Builder::with_capacity(rows))),
            DataType::UInt32 => Ok(Self::UInt32(UInt32Builder::with_capacity(rows))),
            DataType::UInt64 => Ok(Self::UInt64(UInt64Builder::with_capacity(rows))),
            DataType::Float32 => Ok(Self::Float32(Float32Builder::with_capacity(rows))),
            DataType::Float64 => Ok(Self::Float64(Float64Builder::with_capacity(rows))),
            DataType::Utf8 => Ok(Self::Utf8(StringBuilder::with_capacity(rows, rows * 8))),
            DataType::LargeUtf8 => Ok(Self::LargeUtf8(LargeStringBuilder::with_capacity(
                rows,
                rows * 8,
            ))),
            other => Err(DataFusionError::Plan(format!(
                "CDC SQL adapter does not support Arrow field {:?} with type {other:?}",
                field.name()
            ))),
        }
    }

    fn append(&mut self, field: &Field, value: &ColumnValue) -> Result<()> {
        match value {
            ColumnValue::Null => self.append_null(field),
            ColumnValue::ToastUnchanged => Err(DataFusionError::Plan(format!(
                "CDC column {:?} is TOAST-unchanged; include a full before/after value or omit it from the SQL schema",
                field.name()
            ))),
            ColumnValue::Binary(_) => Err(DataFusionError::Plan(format!(
                "CDC column {:?} arrived in binary pgoutput format, but datum-sql CDC decoding expects text values",
                field.name()
            ))),
            ColumnValue::Text(text) => self.append_text(field, text),
        }
    }

    fn append_null(&mut self, field: &Field) -> Result<()> {
        if !field.is_nullable() {
            return Err(DataFusionError::Plan(format!(
                "CDC column {:?} is NULL but the Arrow field is non-nullable",
                field.name()
            )));
        }
        match self {
            Self::Boolean(builder) => builder.append_null(),
            Self::Int8(builder) => builder.append_null(),
            Self::Int16(builder) => builder.append_null(),
            Self::Int32(builder) => builder.append_null(),
            Self::Int64(builder) => builder.append_null(),
            Self::UInt8(builder) => builder.append_null(),
            Self::UInt16(builder) => builder.append_null(),
            Self::UInt32(builder) => builder.append_null(),
            Self::UInt64(builder) => builder.append_null(),
            Self::Float32(builder) => builder.append_null(),
            Self::Float64(builder) => builder.append_null(),
            Self::Utf8(builder) => builder.append_null(),
            Self::LargeUtf8(builder) => builder.append_null(),
        }
        Ok(())
    }

    fn append_text(&mut self, field: &Field, text: &str) -> Result<()> {
        match self {
            Self::Boolean(builder) => builder.append_value(parse_text(field, text)?),
            Self::Int8(builder) => builder.append_value(parse_text(field, text)?),
            Self::Int16(builder) => builder.append_value(parse_text(field, text)?),
            Self::Int32(builder) => builder.append_value(parse_text(field, text)?),
            Self::Int64(builder) => builder.append_value(parse_text(field, text)?),
            Self::UInt8(builder) => builder.append_value(parse_text(field, text)?),
            Self::UInt16(builder) => builder.append_value(parse_text(field, text)?),
            Self::UInt32(builder) => builder.append_value(parse_text(field, text)?),
            Self::UInt64(builder) => builder.append_value(parse_text(field, text)?),
            Self::Float32(builder) => builder.append_value(parse_text(field, text)?),
            Self::Float64(builder) => builder.append_value(parse_text(field, text)?),
            Self::Utf8(builder) => builder.append_value(text),
            Self::LargeUtf8(builder) => builder.append_value(text),
        }
        Ok(())
    }

    fn finish(mut self) -> ArrayRef {
        match &mut self {
            Self::Boolean(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::Int8(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::Int16(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::Int32(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::Int64(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::UInt8(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::UInt16(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::UInt32(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::UInt64(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::Float32(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::Float64(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::Utf8(builder) => Arc::new(builder.finish()) as ArrayRef,
            Self::LargeUtf8(builder) => Arc::new(builder.finish()) as ArrayRef,
        }
    }
}

fn parse_text<T>(field: &Field, text: &str) -> Result<T>
where
    T: std::str::FromStr,
    T::Err: std::fmt::Display,
{
    text.parse::<T>().map_err(|error| {
        DataFusionError::Plan(format!(
            "CDC column {:?} value {:?} cannot be parsed as {:?}: {error}",
            field.name(),
            text,
            field.data_type()
        ))
    })
}