Add routing

Author: fpacifici

sentry_streams/src/consumer.rs

@@ -8,6 +8,8 @@
 use crate::kafka_config::PyKafkaConsumerConfig;
 use crate::operators::build;
 use crate::operators::RuntimeOperator;
+use crate::routes::Route;
+use crate::routes::RoutedValue;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use sentry_arroyo::backends::kafka::types::KafkaPayload;
@@ -39,6 +41,8 @@ pub struct ArroyoConsumer {
 
     topic: String,
 
+    source: String,
+
     steps: Vec<Py<RuntimeOperator>>,
 
     /// The ProcessorHandle allows the main thread to stop the StreamingProcessor
@@ -49,10 +53,11 @@ pub struct ArroyoConsumer {
 #[pymethods]
 impl ArroyoConsumer {
     #[new]
-    fn new(kafka_config: PyKafkaConsumerConfig, topic: String) -> Self {
+    fn new(source: String, kafka_config: PyKafkaConsumerConfig, topic: String) -> Self {
         ArroyoConsumer {
             consumer_config: kafka_config,
             topic,
+            source,
             steps: Vec::new(),
             handle: None,
         }
@@ -77,7 +82,10 @@ impl ArroyoConsumer {
                 .map(|step| step.clone_ref(py))
                 .collect::<Vec<_>>()
         });
-        let factory = ArroyoStreamingFactory { steps: steps_copy };
+        let factory = ArroyoStreamingFactory {
+            source: self.source.clone(),
+            steps: steps_copy,
+        };
         let config = self.consumer_config.clone().into();
         let processor = StreamProcessor::with_kafka(config, factory, Topic::new(&self.topic), None);
         self.handle = Some(processor.get_handle());
@@ -89,7 +97,7 @@ impl ArroyoConsumer {
         })
         .expect("Error setting Ctrl+C handler");
 
-        processor.run();
+        let _ = processor.run();
     }
 
     fn shutdown(&mut self) {
@@ -103,40 +111,45 @@ impl ArroyoConsumer {
 /// Converts a Message<KafkaPayload> to a Message<Py<PyAny>>.
 /// It takes the Kafka payload as bytes and turns it into a
 /// Python bytes object.
-fn to_python(message: Message<KafkaPayload>) -> Message<Py<PyAny>> {
-    let payload = Python::with_gil(|py| {
-        let payload = message.payload().payload().unwrap();
-        let py_bytes = PyBytes::new(py, payload);
-        py_bytes.into_any().unbind()
+fn to_python(source: &str, message: Message<KafkaPayload>) -> Message<RoutedValue> {
+    let payload = Python::with_gil(|py| match message.payload().payload() {
+        Some(payload) => PyBytes::new(py, payload).into_any().unbind(),
+        None => py.None(),
     });
-    message.replace(payload)
+    let route = Route::new(source.to_string(), vec![]);
+    message.replace(RoutedValue { route, payload })
 }
 
 /// Builds the Arroyo StreamProcessor for this consumer.
 /// It plugs a Commit policy at the end and a translator at the beginning
 /// that takes the payload of the Kafka message and turns it into a Py<PyAny>
-fn build_chain(steps: &[Py<RuntimeOperator>]) -> Box<dyn ProcessingStrategy<KafkaPayload>> {
-    let mut next: Box<dyn ProcessingStrategy<Py<PyAny>>> =
+fn build_chain(
+    source: &str,
+    steps: &[Py<RuntimeOperator>],
+) -> Box<dyn ProcessingStrategy<KafkaPayload>> {
+    let mut next: Box<dyn ProcessingStrategy<RoutedValue>> =
         Box::new(CommitOffsets::new(Duration::from_secs(5)));
     for step in steps.iter().rev() {
         next = build(step, next);
     }
 
-    let converter = RunTask::new(
-        |message: Message<KafkaPayload>| Ok(to_python(message)),
-        next,
-    );
+    let copied_source = source.to_string();
+    let conversion_function =
+        move |message: Message<KafkaPayload>| Ok(to_python(&copied_source, message));
+
+    let converter = RunTask::new(conversion_function, next);
 
     Box::new(converter)
 }
 
 struct ArroyoStreamingFactory {
+    source: String,
     steps: Vec<Py<RuntimeOperator>>,
 }
 
 impl ProcessingStrategyFactory<KafkaPayload> for ArroyoStreamingFactory {
     fn create(&self) -> Box<dyn ProcessingStrategy<KafkaPayload>> {
-        build_chain(&self.steps)
+        build_chain(&self.source, &self.steps)
     }
 }
 #[cfg(test)]
@@ -167,6 +180,20 @@ mod tests {
         });
     }
 
+    #[test]
+    fn test_to_none_python() {
+        Python::with_gil(|py| {
+            let message =
+                Message::new_any_message(KafkaPayload::new(None, None, None), BTreeMap::new());
+
+            let python_message = to_python(message);
+
+            let py_payload = python_message.payload();
+
+            assert!(py_payload.is_none(py));
+        });
+    }
+
     #[test]
     fn test_build_chain() {
         Python::with_gil(|py| {

sentry_streams/src/fake_strategy.rs

@@ -1,4 +1,5 @@
 use super::*;
+use crate::routes::RoutedValue;
 use sentry_arroyo::processing::strategies::{
     CommitRequest, ProcessingStrategy, StrategyError, SubmitError,
 };
@@ -10,13 +11,16 @@ pub struct FakeStrategy {
     pub submitted: Arc<Mutex<Vec<Py<PyAny>>>>,
 }
 
-impl ProcessingStrategy<Py<PyAny>> for FakeStrategy {
+impl ProcessingStrategy<RoutedValue> for FakeStrategy {
     fn poll(&mut self) -> Result<Option<CommitRequest>, StrategyError> {
         Ok(None)
     }
 
-    fn submit(&mut self, message: Message<Py<PyAny>>) -> Result<(), SubmitError<Py<PyAny>>> {
-        self.submitted.lock().unwrap().push(message.into_payload());
+    fn submit(&mut self, message: Message<RoutedValue>) -> Result<(), SubmitError<RoutedValue>> {
+        self.submitted
+            .lock()
+            .unwrap()
+            .push(message.into_payload().payload);
         Ok(())
     }
 

sentry_streams/src/operators.rs

@@ -1,5 +1,5 @@
 use crate::kafka_config::PyKafkaConsumerConfig;
-use crate::routes::Route;
+use crate::routes::{Route, RoutedValue};
 use crate::transformer::build_map;
 use pyo3::prelude::*;
 use sentry_arroyo::processing::strategies::ProcessingStrategy;
@@ -77,8 +77,8 @@ impl RuntimeOperator {
 
 pub fn build(
     step: &Py<RuntimeOperator>,
-    next: Box<dyn ProcessingStrategy<Py<PyAny>>>,
-) -> Box<dyn ProcessingStrategy<Py<PyAny>>> {
+    next: Box<dyn ProcessingStrategy<RoutedValue>>,
+) -> Box<dyn ProcessingStrategy<RoutedValue>> {
     match step.get() {
         RuntimeOperator::Map { function, .. } => {
             let func_ref = Python::with_gil(|py| function.clone_ref(py));

sentry_streams/src/transformer.rs

@@ -1,14 +1,18 @@
+use crate::routes::RoutedValue;
 use pyo3::prelude::*;
 use sentry_arroyo::processing::strategies::run_task::RunTask;
 use sentry_arroyo::processing::strategies::ProcessingStrategy;
 use sentry_arroyo::types::Message;
 
 /// Executes a Python callable with an Arroyo message containing Any and
 /// returns the result.
-fn call_python_function(callable: &Py<PyAny>, message: &Message<Py<PyAny>>) -> Py<PyAny> {
+fn call_python_function(
+    callable: &Py<PyAny>,
+    message: &Message<RoutedValue>,
+) -> Result<Py<PyAny>, PyErr> {
     Python::with_gil(|py| {
-        let result = callable.call1(py, (message.payload(),)).unwrap();
-        result
+        let python_payload = message.payload().payload.clone_ref(py);
+        callable.call1(py, (python_payload,))
     })
 }
 
@@ -20,11 +24,17 @@ fn call_python_function(callable: &Py<PyAny>, message: &Message<Py<PyAny>>) -> P
 /// This function takes a `next`  step to wire the Arroyo strategy to.
 pub fn build_map(
     callable: Py<PyAny>,
-    next: Box<dyn ProcessingStrategy<Py<PyAny>>>,
-) -> Box<dyn ProcessingStrategy<Py<PyAny>>> {
-    let mapper = move |message| {
+    next: Box<dyn ProcessingStrategy<RoutedValue>>,
+) -> Box<dyn ProcessingStrategy<RoutedValue>> {
+    let mapper = move |message: Message<RoutedValue>| {
+        let route = message.payload().route.clone();
         let transformed = call_python_function(&callable, &message);
-        Ok(message.replace(transformed))
+        // TODO: Create an exception for Invalid messages in Python
+        // This now crashes if the Python code fails.
+        Ok(message.replace(RoutedValue {
+            route,
+            payload: transformed.unwrap(),
+        }))
     };
     Box::new(RunTask::new(mapper, next))
 }
@@ -33,13 +43,23 @@ mod tests {
     use super::*;
     use crate::fake_strategy::assert_messages_match;
     use crate::fake_strategy::FakeStrategy;
+    use crate::routes::Route;
     use pyo3::ffi::c_str;
     use pyo3::IntoPyObjectExt;
     use sentry_arroyo::processing::strategies::ProcessingStrategy;
     use std::collections::BTreeMap;
     use std::ops::Deref;
     use std::sync::{Arc, Mutex};
 
+    fn build_routed_value(py: Python<'_>) -> RoutedValue {
+        let msg_payload = "test_message".into_py_any(py).unwrap();
+        let route = Route::new("source".to_string(), vec!["waypoint".to_string()]);
+        RoutedValue {
+            route,
+            payload: msg_payload,
+        }
+    }
+
     #[test]
     fn test_build_map() {
         Python::with_gil(|py| {
@@ -56,8 +76,7 @@ mod tests {
 
             let mut strategy = build_map(callable, Box::new(next_step));
 
-            let message =
-                Message::new_any_message("test_message".into_py_any(py).unwrap(), BTreeMap::new());
+            let message = Message::new_any_message(build_routed_value(py), BTreeMap::new());
             let result = strategy.submit(message);
 
             assert!(result.is_ok());
@@ -79,10 +98,9 @@ mod tests {
                 .into_py_any(py)
                 .unwrap();
 
-            let message =
-                Message::new_any_message("test_message".into_py_any(py).unwrap(), BTreeMap::new());
+            let message = Message::new_any_message(build_routed_value(py), BTreeMap::new());
 
-            let result = call_python_function(&callable, &message);
+            let result = call_python_function(&callable, &message).unwrap();
 
             assert_eq!(
                 result.extract::<String>(py).unwrap(),