use crate::{
ConnectorConfig, ConnectorError, ConnectorResult, ConsumerConfig, ProducerConfig, SinkRecord,
SourceRecord,
};
use async_trait::async_trait;
use tokio::sync::mpsc;
#[async_trait]
pub trait SinkConnector: Send + Sync {
async fn initialize(&mut self, config: ConnectorConfig) -> ConnectorResult<()>;
async fn consumer_configs(&self) -> ConnectorResult<Vec<ConsumerConfig>>;
async fn process_batch(&mut self, records: Vec<SinkRecord>) -> ConnectorResult<()>;
async fn shutdown(&mut self) -> ConnectorResult<()> {
Ok(())
}
async fn health_check(&self) -> ConnectorResult<()> {
Ok(())
}
}
#[async_trait]
pub trait SourceConnector: Send + Sync {
async fn initialize(&mut self, config: ConnectorConfig) -> ConnectorResult<()>;
async fn producer_configs(&self) -> ConnectorResult<Vec<ProducerConfig>>;
fn mode(&self) -> SourceConnectorMode {
SourceConnectorMode::Polling
}
async fn start_streaming(&mut self, _sender: SourceSender) -> ConnectorResult<()> {
Err(ConnectorError::config(
"start_streaming() not implemented for this source connector",
))
}
async fn poll(&mut self) -> ConnectorResult<Vec<SourceEnvelope>> {
Err(ConnectorError::config(
"poll() not implemented for this source connector",
))
}
async fn commit(&mut self, offsets: Vec<Offset>) -> ConnectorResult<()> {
let _ = offsets; Ok(())
}
async fn shutdown(&mut self) -> ConnectorResult<()> {
Ok(())
}
async fn health_check(&self) -> ConnectorResult<()> {
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SourceConnectorMode {
Polling,
Streaming,
}
#[derive(Clone)]
pub struct SourceSender {
sender: mpsc::Sender<SourceEnvelope>,
}
impl SourceSender {
pub(crate) fn new(sender: mpsc::Sender<SourceEnvelope>) -> Self {
Self { sender }
}
pub async fn send(&self, envelope: impl Into<SourceEnvelope>) -> ConnectorResult<()> {
self.sender
.send(envelope.into())
.await
.map_err(|e| ConnectorError::fatal(format!("failed to emit source envelope: {}", e)))
}
pub async fn send_batch<I>(&self, envelopes: I) -> ConnectorResult<()>
where
I: IntoIterator,
I::Item: Into<SourceEnvelope>,
{
for envelope in envelopes {
self.send(envelope).await?;
}
Ok(())
}
pub fn is_closed(&self) -> bool {
self.sender.is_closed()
}
}
#[derive(Debug, Clone)]
pub struct SourceEnvelope {
pub record: SourceRecord,
pub offset: Option<Offset>,
}
impl SourceEnvelope {
pub fn new(record: SourceRecord) -> Self {
Self {
record,
offset: None,
}
}
pub fn with_offset(record: SourceRecord, offset: Offset) -> Self {
Self {
record,
offset: Some(offset),
}
}
pub fn record(&self) -> &SourceRecord {
&self.record
}
pub fn offset(&self) -> Option<&Offset> {
self.offset.as_ref()
}
pub fn into_parts(self) -> (SourceRecord, Option<Offset>) {
(self.record, self.offset)
}
}
impl From<SourceRecord> for SourceEnvelope {
fn from(record: SourceRecord) -> Self {
Self::new(record)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Offset {
pub partition: String,
pub value: u64,
pub metadata: Option<String>,
}
impl Offset {
pub fn new(partition: impl Into<String>, value: u64) -> Self {
Self {
partition: partition.into(),
value,
metadata: None,
}
}
pub fn with_metadata(
partition: impl Into<String>,
value: u64,
metadata: impl Into<String>,
) -> Self {
Self {
partition: partition.into(),
value,
metadata: Some(metadata.into()),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_offset_creation() {
let offset = Offset::new("partition-0", 42);
assert_eq!(offset.partition, "partition-0");
assert_eq!(offset.value, 42);
assert!(offset.metadata.is_none());
let offset_with_meta = Offset::with_metadata("partition-1", 100, "some-metadata");
assert_eq!(offset_with_meta.metadata, Some("some-metadata".to_string()));
}
}