datum_sql/

sink.rs

use std::fmt;
use std::sync::Arc;

use arrow::record_batch::RecordBatch;
use datafusion::common::{DataFusionError, Result};
use datum::{Source, StreamError, StreamResult};

use crate::connect::BatchEnvelope;
use crate::{ChangelogBatch, SqlEventPayload, stream_error};

/// Type-erased source position carried by SQL sink pipelines.
///
/// Connector adapters keep their precise position types at the edge and convert
/// them into this enum when a query needs commit-after-sink-confirmation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SqlSourcePosition {
    #[cfg(feature = "mq")]
    /// Kafka topic/partition offset ranges.
    Kafka(crate::connect::mq::KafkaSourcePosition),
    #[cfg(feature = "cdc")]
    /// PostgreSQL CDC LSN position.
    Cdc(crate::connect::cdc::CdcSourcePosition),
    /// Opaque source position for user-defined connector adapters.
    Custom {
        /// Connector family or instance name.
        connector: Arc<str>,
        /// Connector-specific position text.
        position: Arc<str>,
    },
}

impl SqlSourcePosition {
    /// Creates an opaque connector-specific source position.
    #[must_use]
    pub fn custom(connector: impl Into<Arc<str>>, position: impl Into<Arc<str>>) -> Self {
        Self::Custom {
            connector: connector.into(),
            position: position.into(),
        }
    }
}

type CommitCallback = dyn Fn() -> StreamResult<()> + Send + Sync;

/// A committable source position.
///
/// SQL sink lowering calls this only after the sink has confirmed the batch.
#[derive(Clone)]
pub struct SourceCommit {
    description: Arc<str>,
    commit: Arc<CommitCallback>,
}

impl SourceCommit {
    /// Creates a commit handle from a callback.
    #[must_use]
    pub fn from_fn<F>(description: impl Into<Arc<str>>, commit: F) -> Self
    where
        F: Fn() -> StreamResult<()> + Send + Sync + 'static,
    {
        Self {
            description: description.into(),
            commit: Arc::new(commit),
        }
    }

    /// Creates a commit handle that does nothing.
    #[must_use]
    pub fn noop() -> Self {
        Self::from_fn("no source position", || Ok(()))
    }

    /// Confirms the source position represented by this handle.
    pub fn commit(&self) -> StreamResult<()> {
        (self.commit)()
    }

    /// Human-readable description used in debug output.
    #[must_use]
    pub fn description(&self) -> &str {
        &self.description
    }
}

impl fmt::Debug for SourceCommit {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SourceCommit")
            .field("description", &self.description)
            .finish_non_exhaustive()
    }
}

/// Append-only SQL batch plus source metadata and a deferred commit handle.
#[derive(Clone)]
pub struct CommittableRecordBatch {
    batch: RecordBatch,
    envelope: BatchEnvelope<Option<SqlSourcePosition>>,
    commit: SourceCommit,
}

impl CommittableRecordBatch {
    /// Creates an append-only SQL batch with source metadata and a commit hook.
    #[must_use]
    pub fn new(
        batch: RecordBatch,
        source_position: Option<SqlSourcePosition>,
        schema_revision: u64,
        commit: SourceCommit,
    ) -> Self {
        Self {
            batch,
            envelope: BatchEnvelope::new(source_position, schema_revision),
            commit,
        }
    }

    /// Wraps an ordinary Arrow batch with no source position.
    #[must_use]
    pub fn plain(batch: RecordBatch) -> Self {
        Self::new(batch, None, 0, SourceCommit::noop())
    }

    /// Returns the Arrow payload batch.
    #[must_use]
    pub fn batch(&self) -> &RecordBatch {
        &self.batch
    }

    /// Returns the connector-side envelope.
    #[must_use]
    pub fn envelope(&self) -> &BatchEnvelope<Option<SqlSourcePosition>> {
        &self.envelope
    }

    /// Returns the optional source position.
    #[must_use]
    pub fn source_position(&self) -> Option<&SqlSourcePosition> {
        self.envelope.source_position().as_ref()
    }

    /// Returns the schema revision attached by the source adapter.
    #[must_use]
    pub fn schema_revision(&self) -> u64 {
        self.envelope.schema_revision()
    }

    /// Returns the deferred source commit handle.
    #[must_use]
    pub fn commit_handle(&self) -> &SourceCommit {
        &self.commit
    }

    /// Commits the source position after a sink write has succeeded.
    pub fn commit(&self) -> StreamResult<()> {
        self.commit.commit()
    }

    pub(crate) fn try_map_batch<F>(self, f: F) -> Result<Self>
    where
        F: FnOnce(&RecordBatch) -> Result<RecordBatch>,
    {
        let batch = f(&self.batch)?;
        Ok(Self {
            batch,
            envelope: self.envelope,
            commit: self.commit,
        })
    }
}

impl fmt::Debug for CommittableRecordBatch {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("CommittableRecordBatch")
            .field("batch", &self.batch)
            .field("envelope", &self.envelope)
            .field("commit", &self.commit)
            .finish()
    }
}

impl SqlEventPayload for CommittableRecordBatch {
    fn event_time_batch(&self) -> &RecordBatch {
        &self.batch
    }

    fn event_time_partition(&self, row: usize) -> Option<i64> {
        match self.source_position()? {
            #[cfg(feature = "mq")]
            SqlSourcePosition::Kafka(position) => position.partition_for_row(row).map(i64::from),
            #[cfg(feature = "cdc")]
            SqlSourcePosition::Cdc(_) | SqlSourcePosition::Custom { .. } => None,
            #[cfg(not(feature = "cdc"))]
            SqlSourcePosition::Custom { .. } => None,
        }
    }

    fn event_time_active_partitions(&self) -> Option<Vec<i64>> {
        match self.source_position()? {
            #[cfg(feature = "mq")]
            SqlSourcePosition::Kafka(position) => {
                if !position.has_event_time_partition_mapping() {
                    return None;
                }
                let partitions = position
                    .active_partitions()
                    .iter()
                    .map(|partition| i64::from(partition.partition))
                    .collect::<Vec<_>>();
                (!partitions.is_empty()).then_some(partitions)
            }
            #[cfg(feature = "cdc")]
            SqlSourcePosition::Cdc(_) | SqlSourcePosition::Custom { .. } => None,
            #[cfg(not(feature = "cdc"))]
            SqlSourcePosition::Custom { .. } => None,
        }
    }
}

type AppendWriter = dyn Fn(RecordBatch) -> StreamResult<()> + Send + Sync;
type ChangelogWriter = dyn Fn(ChangelogBatch) -> StreamResult<()> + Send + Sync;

#[derive(Clone)]
pub(crate) struct RegisteredSqlSink {
    append_writer: Arc<AppendWriter>,
    changelog_writer: Option<Arc<ChangelogWriter>>,
}

impl RegisteredSqlSink {
    pub(crate) fn append_only<F>(writer: F) -> Self
    where
        F: Fn(RecordBatch) -> StreamResult<()> + Send + Sync + 'static,
    {
        Self {
            append_writer: Arc::new(writer),
            changelog_writer: None,
        }
    }

    pub(crate) fn changelog<F, G>(append_writer: F, changelog_writer: G) -> Self
    where
        F: Fn(RecordBatch) -> StreamResult<()> + Send + Sync + 'static,
        G: Fn(ChangelogBatch) -> StreamResult<()> + Send + Sync + 'static,
    {
        Self {
            append_writer: Arc::new(append_writer),
            changelog_writer: Some(Arc::new(changelog_writer)),
        }
    }

    #[cfg(feature = "mq")]
    pub(crate) fn kafka_json(
        settings: datum_mq::KafkaProducerSettings,
        topic: String,
        format: crate::connect::mq::JsonRowFormat,
    ) -> Self {
        Self::append_only(move |batch| {
            let payloads = format.encode_record_batch(&batch).map_err(stream_error)?;
            if payloads.is_empty() {
                return Ok(());
            }

            let records = payloads.into_iter().map(|payload| {
                datum_mq::ProducerRecord::new(topic.clone(), bytes::Bytes::from(payload))
            });
            let control = Source::from_iter(records)
                .run_with(datum_mq::KafkaSink::plain(settings.clone()))
                .map_err(stream_error)?;
            control.drain_and_shutdown().map_err(stream_error)
        })
    }

    #[must_use]
    pub(crate) fn accepts_changelog(&self) -> bool {
        self.changelog_writer.is_some()
    }

    pub(crate) fn write_append(&self, batch: RecordBatch) -> StreamResult<()> {
        (self.append_writer)(batch)
    }

    pub(crate) fn write_changelog(&self, changes: ChangelogBatch) -> StreamResult<()> {
        let Some(writer) = &self.changelog_writer else {
            return Err(StreamError::Failed(
                "datum-sql sink is not changelog-aware".to_owned(),
            ));
        };
        writer(changes)
    }
}

impl fmt::Debug for RegisteredSqlSink {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RegisteredSqlSink")
            .field("accepts_changelog", &self.accepts_changelog())
            .finish_non_exhaustive()
    }
}

pub(crate) enum AppendSinkSource {
    Plain(Source<RecordBatch>),
    Committable(Source<CommittableRecordBatch>),
}

impl AppendSinkSource {
    pub(crate) fn into_committable(self) -> Source<CommittableRecordBatch> {
        match self {
            Self::Plain(source) => source.map(CommittableRecordBatch::plain),
            Self::Committable(source) => source,
        }
    }
}

pub(crate) fn sink_registry_error(error: impl fmt::Display) -> DataFusionError {
    DataFusionError::External(Box::new(SinkRegistryError(error.to_string())))
}

#[derive(Debug)]
struct SinkRegistryError(String);

impl fmt::Display for SinkRegistryError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.0)
    }
}

impl std::error::Error for SinkRegistryError {}