orion-server 0.2.0

Declarative services runtime powered by dataflow-rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use std::sync::Arc;

use async_trait::async_trait;
use dataflow_rs::engine::error::DataflowError;
use dataflow_rs::engine::functions::AsyncFunctionHandler;
use dataflow_rs::engine::functions::PublishKafkaConfig;
use dataflow_rs::engine::task_context::TaskContext;
use dataflow_rs::engine::task_outcome::TaskOutcome;
use serde_json::Value;

use crate::connector::ConnectorRegistry;

/// Kafka publish handler.
pub struct PublishKafkaHandler {
    pub registry: Arc<ConnectorRegistry>,
    pub producer: Option<Arc<crate::kafka::producer::KafkaProducer>>,
}

#[async_trait]
impl AsyncFunctionHandler for PublishKafkaHandler {
    type Input = PublishKafkaConfig;

    async fn execute(
        &self,
        ctx: &mut TaskContext<'_>,
        input: &PublishKafkaConfig,
    ) -> dataflow_rs::Result<TaskOutcome> {
        crate::engine::profile::record("publish_kafka", Some(&input.connector), async move {
            // Verify connector exists.
            let _connector = self.registry.get(&input.connector).await.ok_or_else(|| {
                DataflowError::function_execution(
                    format!("Connector '{}' not found", input.connector),
                    None,
                )
            })?;

            let producer = match &self.producer {
                Some(p) => p,
                None => {
                    return Err(DataflowError::FunctionExecution {
                        context: format!(
                            "Kafka publishing to topic '{}' is not available. \
                         Enable Kafka in configuration to use publish_kafka.",
                            input.topic
                        ),
                        source: None,
                    });
                }
            };

            let key = if let Some(compiled) = input.compiled_key_logic.as_deref() {
                let result: Value = ctx
                    .datalogic()
                    .session()
                    .eval_into(compiled, &ctx.message().context)
                    .map_err(|e| DataflowError::LogicEvaluation(e.to_string()))?;
                let key_str = if let Some(s) = result.as_str() {
                    s.to_string()
                } else {
                    serde_json::to_string(&result).map_err(|e| {
                        DataflowError::function_execution(
                            format!("Failed to serialize Kafka message key: {e}"),
                            None,
                        )
                    })?
                };
                Some(key_str)
            } else {
                None
            };

            let value = if let Some(compiled) = input.compiled_value_logic.as_deref() {
                let result: Value = ctx
                    .datalogic()
                    .session()
                    .eval_into(compiled, &ctx.message().context)
                    .map_err(|e| DataflowError::LogicEvaluation(e.to_string()))?;
                serde_json::to_string(&result).map_err(|e| {
                    DataflowError::function_execution(
                        format!("Failed to serialize Kafka message value: {e}"),
                        None,
                    )
                })?
            } else {
                let data_json: Value = ctx.data().into();
                serde_json::to_string(&data_json).map_err(|e| {
                    DataflowError::function_execution(
                        format!("Failed to serialize Kafka message value: {e}"),
                        None,
                    )
                })?
            };

            producer
                .send(&input.topic, key.as_deref(), value.as_bytes())
                .await
                .map_err(|e| {
                    DataflowError::function_execution(
                        format!("Kafka publish to '{}' failed: {e}", input.topic),
                        None,
                    )
                })?;

            tracing::debug!(
                topic = %input.topic,
                "Published message to Kafka"
            );

            Ok(TaskOutcome::Success)
        })
        .await
    }
}