dora_operator_api_python/
lib.rs

1use std::{
2    collections::{BTreeMap, HashMap},
3    sync::{Arc, Mutex},
4};
5
6use arrow::pyarrow::ToPyArrow;
7use chrono::{DateTime, Utc};
8use dora_node_api::{
9    DoraNode, Event, EventStream, Metadata, MetadataParameters, Parameter, StopCause,
10    merged::{MergeExternalSend, MergedEvent},
11};
12use eyre::{Context, Result};
13use futures::{Stream, StreamExt};
14use futures_concurrency::stream::Merge as _;
15use pyo3::{
16    prelude::*,
17    sync::PyOnceLock,
18    types::{IntoPyDict, PyBool, PyDict, PyFloat, PyInt, PyList, PyModule, PyString, PyTuple},
19};
20use std::time::UNIX_EPOCH;
21
22/// Cached Python `datetime` module to avoid repeated `PyModule::import` on the hot path.
23static DATETIME_MODULE: PyOnceLock<Py<PyModule>> = PyOnceLock::new();
24
25pub fn datetime_module<'py>(py: Python<'py>) -> PyResult<&'py Bound<'py, PyModule>> {
26    Ok(DATETIME_MODULE
27        .get_or_try_init(py, || PyModule::import(py, "datetime").map(|m| m.unbind()))?
28        .bind(py))
29}
30
31/// Dora Event
32pub struct PyEvent {
33    pub event: MergedEvent<Py<PyAny>>,
34}
35
36/// Keeps the dora node alive until all event objects have been dropped.
37#[derive(Clone)]
38#[pyclass(skip_from_py_object)]
39pub struct NodeCleanupHandle {
40    pub _handles: Arc<CleanupHandle<DoraNode>>,
41}
42
43/// Owned type with delayed cleanup (using `handle` method).
44pub struct DelayedCleanup<T>(Arc<Mutex<T>>);
45
46impl<T> Clone for DelayedCleanup<T> {
47    fn clone(&self) -> Self {
48        Self(self.0.clone())
49    }
50}
51
52impl<T> DelayedCleanup<T> {
53    pub fn new(value: T) -> Self {
54        Self(Arc::new(Mutex::new(value)))
55    }
56
57    pub fn handle(&self) -> CleanupHandle<T> {
58        CleanupHandle(self.0.clone())
59    }
60
61    pub fn get_mut(&self) -> std::sync::MutexGuard<'_, T> {
62        self.0.try_lock().expect("failed to lock DelayedCleanup")
63    }
64}
65
66impl Stream for DelayedCleanup<EventStream> {
67    type Item = Event;
68
69    fn poll_next(
70        self: std::pin::Pin<&mut Self>,
71        cx: &mut std::task::Context<'_>,
72    ) -> std::task::Poll<Option<Self::Item>> {
73        let mut inner: std::sync::MutexGuard<'_, EventStream> = self.get_mut().get_mut();
74        inner.poll_next_unpin(cx)
75    }
76}
77
78impl<'a, E> MergeExternalSend<'a, E> for DelayedCleanup<EventStream>
79where
80    E: 'static,
81{
82    type Item = MergedEvent<E>;
83
84    fn merge_external_send(
85        self,
86        external_events: impl Stream<Item = E> + Unpin + Send + Sync + 'a,
87    ) -> Box<dyn Stream<Item = Self::Item> + Unpin + Send + Sync + 'a> {
88        let dora = self.map(MergedEvent::Dora);
89        let external = external_events.map(MergedEvent::External);
90        Box::new((dora, external).merge())
91    }
92}
93
94#[allow(dead_code)]
95pub struct CleanupHandle<T>(Arc<Mutex<T>>);
96
97impl PyEvent {
98    pub fn to_py_dict(self, py: Python<'_>) -> PyResult<Py<PyDict>> {
99        let mut pydict = HashMap::new();
100        match &self.event {
101            MergedEvent::Dora(_) => pydict.insert(
102                "kind",
103                "dora"
104                    .into_pyobject(py)
105                    .context("Failed to create pystring")?
106                    .unbind()
107                    .into(),
108            ),
109            MergedEvent::External(_) => pydict.insert(
110                "kind",
111                "external"
112                    .into_pyobject(py)
113                    .context("Failed to create pystring")?
114                    .unbind()
115                    .into(),
116            ),
117        };
118        match &self.event {
119            MergedEvent::Dora(event) => {
120                if let Some(id) = Self::id(event) {
121                    pydict.insert(
122                        "id",
123                        id.into_pyobject(py)
124                            .context("Failed to create id pyobject")?
125                            .into(),
126                    );
127                }
128                pydict.insert(
129                    "type",
130                    Self::ty(event)
131                        .into_pyobject(py)
132                        .context("Failed to create event pyobject")?
133                        .unbind()
134                        .into(),
135                );
136
137                if let Some(value) = self.value(py)? {
138                    pydict.insert("value", value);
139                }
140                if let Some(metadata) = Self::metadata(event, py)? {
141                    pydict.insert("metadata", metadata);
142                }
143                if let Some(error) = Self::error(event) {
144                    pydict.insert(
145                        "error",
146                        error
147                            .into_pyobject(py)
148                            .context("Failed to create error pyobject")?
149                            .unbind()
150                            .into(),
151                    );
152                }
153            }
154            MergedEvent::External(event) => {
155                pydict.insert("value", event.clone_ref(py));
156            }
157        }
158
159        Ok(pydict
160            .into_py_dict(py)
161            .context("Failed to create py_dict")?
162            .unbind())
163    }
164
165    fn ty(event: &Event) -> &'static str {
166        match event {
167            Event::Stop(_) => "STOP",
168            Event::Input { .. } => "INPUT",
169            Event::InputClosed { .. } => "INPUT_CLOSED",
170            Event::InputRecovered { .. } => "INPUT_RECOVERED",
171            Event::NodeRestarted { .. } => "NODE_RESTARTED",
172            Event::Reload { .. } => "RELOAD",
173            Event::ParamUpdate { .. } => "PARAM_UPDATE",
174            Event::ParamDeleted { .. } => "PARAM_DELETED",
175            Event::NodeFailed { .. } => "NODE_FAILED",
176            Event::Error(_) => "ERROR",
177            // `Event` is `#[non_exhaustive]`; surface genuinely new variants
178            // as UNKNOWN rather than failing to build on future dora upgrades.
179            _other => "UNKNOWN",
180        }
181    }
182
183    fn id(event: &Event) -> Option<&str> {
184        match event {
185            Event::Input { id, .. } => Some(id.as_ref()),
186            Event::InputClosed { id } => Some(id.as_ref()),
187            Event::InputRecovered { id } => Some(id.as_ref()),
188            Event::NodeRestarted { id } => Some(id.as_ref()),
189            Event::Reload { operator_id } => operator_id.as_ref().map(|id| id.as_ref()),
190            Event::ParamUpdate { key, .. } => Some(key.as_str()),
191            Event::ParamDeleted { key } => Some(key.as_str()),
192            Event::NodeFailed { source_node_id, .. } => Some(source_node_id.as_ref()),
193            Event::Stop(cause) => match cause {
194                StopCause::Manual => Some("MANUAL"),
195                StopCause::AllInputsClosed => Some("ALL_INPUTS_CLOSED"),
196                &_ => None,
197            },
198            _ => None,
199        }
200    }
201
202    /// Returns the payload of an event as a Python object (if any).
203    ///
204    /// - `Input`: the Arrow array payload.
205    /// - `ParamUpdate`: the new parameter value, converted from JSON.
206    fn value(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
207        match &self.event {
208            MergedEvent::Dora(Event::Input { data, .. }) => {
209                // Ownership transfer: to_data() clones buffer Arcs, to_pyarrow()
210                // creates an FFI_ArrowArray with a release callback. PyArrow's
211                // _import_from_c takes ownership and calls release when the Python
212                // object is GC'd, which drops the cloned Arcs. No leak.
213                let array_data = data.to_data().to_pyarrow(py)?.unbind();
214                Ok(Some(array_data))
215            }
216            MergedEvent::Dora(Event::ParamUpdate { value, .. }) => {
217                // `pythonize` converts serde_json::Value recursively into native
218                // Python types (None, bool, int, float, str, list, dict).
219                let obj = pythonize::pythonize(py, value)
220                    .context("failed to convert ParamUpdate value to Python")?
221                    .unbind();
222                Ok(Some(obj))
223            }
224            MergedEvent::Dora(Event::NodeFailed {
225                affected_input_ids,
226                error,
227                source_node_id,
228            }) => {
229                let dict = pyo3::types::PyDict::new(py);
230                dict.set_item("source_node_id", source_node_id.as_ref())?;
231                dict.set_item("error", error.as_str())?;
232                let ids: Vec<&str> = affected_input_ids.iter().map(|id| id.as_ref()).collect();
233                dict.set_item("affected_input_ids", ids)?;
234                Ok(Some(dict.unbind().into()))
235            }
236            _ => Ok(None),
237        }
238    }
239
240    fn metadata(event: &Event, py: Python<'_>) -> Result<Option<Py<PyAny>>> {
241        match event {
242            Event::Input { metadata, .. } => Ok(Some(
243                metadata_to_pydict(metadata, py)
244                    .context("Issue deserializing metadata")?
245                    .into_pyobject(py)
246                    .context("Failed to create metadata_to_pydice")?
247                    .unbind()
248                    .into(),
249            )),
250            _ => Ok(None),
251        }
252    }
253
254    fn error(event: &Event) -> Option<&str> {
255        match event {
256            Event::Error(error) => Some(error),
257            _other => None,
258        }
259    }
260}
261
262pub fn pydict_to_metadata(dict: Option<Bound<'_, PyDict>>) -> Result<MetadataParameters> {
263    let mut parameters = BTreeMap::default();
264    if let Some(pymetadata) = dict {
265        for (key, value) in pymetadata.iter() {
266            let key = key.extract::<String>().context("Parsing metadata keys")?;
267            if value.is_exact_instance_of::<PyBool>() {
268                parameters.insert(key, Parameter::Bool(value.extract()?))
269            } else if value.is_instance_of::<PyInt>() {
270                parameters.insert(key, Parameter::Integer(value.extract::<i64>()?))
271            } else if value.is_instance_of::<PyFloat>() {
272                parameters.insert(key, Parameter::Float(value.extract::<f64>()?))
273            } else if value.is_instance_of::<PyString>() {
274                parameters.insert(key, Parameter::String(value.extract()?))
275            } else if (value.is_instance_of::<PyTuple>() || value.is_instance_of::<PyList>())
276                && value.len()? == 0
277            {
278                // Empty list/tuple: no element type to infer. Preserve
279                // list-ness with an empty list (round-trips to `[]`) instead of
280                // coercing to the string "[]" (dora-rs/dora#2027).
281                parameters.insert(key, Parameter::ListString(vec![]))
282            } else if (value.is_instance_of::<PyTuple>() || value.is_instance_of::<PyList>())
283                && value.len()? > 0
284                && value.get_item(0)?.is_exact_instance_of::<PyInt>()
285            {
286                let list: Vec<i64> = value.extract()?;
287                parameters.insert(key, Parameter::ListInt(list))
288            } else if (value.is_instance_of::<PyTuple>() || value.is_instance_of::<PyList>())
289                && value.len()? > 0
290                && value.get_item(0)?.is_exact_instance_of::<PyFloat>()
291            {
292                let list: Vec<f64> = value.extract()?;
293                parameters.insert(key, Parameter::ListFloat(list))
294            } else if (value.is_instance_of::<PyTuple>() || value.is_instance_of::<PyList>())
295                && value.len()? > 0
296                && value.get_item(0)?.is_exact_instance_of::<PyString>()
297            {
298                let list: Vec<String> = value.extract()?;
299                parameters.insert(key, Parameter::ListString(list))
300            } else {
301                // Check if it's a datetime.datetime object
302                let dt_module =
303                    datetime_module(value.py()).context("Failed to import datetime module")?;
304                let datetime_class = dt_module.getattr("datetime")?;
305
306                if value.is_instance(datetime_class.as_ref())? {
307                    // Extract timestamp using timestamp() method
308                    let timestamp_float: f64 = value
309                        .call_method0("timestamp")?
310                        .extract()
311                        .context("Failed to extract timestamp from datetime")?;
312
313                    // Convert to chrono::DateTime<Utc>
314                    // timestamp() returns seconds since epoch as float
315                    // Convert to SystemTime first, then to DateTime<Utc>
316                    let system_time = if timestamp_float >= 0.0 {
317                        let duration = std::time::Duration::try_from_secs_f64(timestamp_float)
318                            .context("Failed to convert timestamp to Duration")?;
319                        UNIX_EPOCH + duration
320                    } else {
321                        let duration = std::time::Duration::try_from_secs_f64(-timestamp_float)
322                            .context("Failed to convert timestamp to Duration")?;
323                        UNIX_EPOCH.checked_sub(duration).unwrap_or(UNIX_EPOCH)
324                    };
325
326                    let dt = DateTime::<Utc>::from(system_time);
327
328                    parameters.insert(key, Parameter::Timestamp(dt))
329                } else {
330                    tracing::warn!(
331                        "unsupported metadata value type for key `{key}`; \
332                         coercing to its string representation: {value}"
333                    );
334                    parameters.insert(key, Parameter::String(value.str()?.to_string()))
335                }
336            };
337        }
338    }
339    Ok(parameters)
340}
341
342pub fn metadata_to_pydict<'a>(
343    metadata: &'a Metadata,
344    py: Python<'a>,
345) -> Result<pyo3::Bound<'a, PyDict>> {
346    let dict = PyDict::new(py);
347
348    // Add timestamp as timezone-aware Python datetime (UTC)
349    // Note: uhlc::Timestamp is a Hybrid Logical Clock. We use get_time().to_system_time()
350    // which extracts the physical clock component. This pattern is used consistently
351    // throughout the dora codebase (e.g., in binaries/daemon/src/log.rs, binaries/coordinator/src/lib.rs)
352    // and assumes the physical time component represents UTC wall-clock time.
353    let timestamp = metadata.timestamp();
354    let system_time = timestamp.get_time().to_system_time();
355    let duration_since_epoch = system_time
356        .duration_since(UNIX_EPOCH)
357        .context("Failed to calculate duration since epoch")?;
358
359    // Extract seconds and microseconds (Python datetime supports microsecond precision)
360    let seconds = duration_since_epoch.as_secs() as i64;
361    let microseconds = duration_since_epoch.subsec_micros();
362
363    // Get UTC timezone from Python's datetime module and create timezone-aware datetime
364    // We use Python's datetime.fromtimestamp() to create a UTC-aware datetime object
365    // This avoids float precision loss by using integer seconds and microseconds
366    let dt_module = datetime_module(py).context("Failed to import datetime module")?;
367    let datetime_class = dt_module.getattr("datetime")?;
368    let utc_timezone = dt_module.getattr("timezone")?.getattr("utc")?;
369
370    // Create timezone-aware datetime using fromtimestamp
371    // We compute total_seconds as float (required by fromtimestamp) but preserve
372    // precision by computing from integer seconds and microseconds separately
373    let total_seconds = seconds as f64 + microseconds as f64 / 1_000_000.0;
374    let py_datetime = datetime_class
375        .call_method1("fromtimestamp", (total_seconds, utc_timezone))
376        .context("Failed to create Python datetime from timestamp")?;
377
378    dict.set_item("timestamp", py_datetime)
379        .context("Could not insert timestamp into python dictionary")?;
380
381    // Add existing parameters
382    for (k, v) in metadata.parameters.iter() {
383        match v {
384            Parameter::Bool(bool) => dict
385                .set_item(k, bool)
386                .context("Could not insert metadata into python dictionary")?,
387            Parameter::Integer(int) => dict
388                .set_item(k, int)
389                .context("Could not insert metadata into python dictionary")?,
390            Parameter::Float(float) => dict
391                .set_item(k, float)
392                .context("Could not insert metadata into python dictionary")?,
393            Parameter::String(s) => dict
394                .set_item(k, s)
395                .context("Could not insert metadata into python dictionary")?,
396            Parameter::ListInt(l) => dict
397                .set_item(k, l)
398                .context("Could not insert metadata into python dictionary")?,
399            Parameter::ListFloat(l) => dict
400                .set_item(k, l)
401                .context("Could not insert metadata into python dictionary")?,
402            Parameter::ListString(l) => dict
403                .set_item(k, l)
404                .context("Could not insert metadata into python dictionary")?,
405            Parameter::Timestamp(dt) => {
406                // Convert chrono::DateTime<Utc> to Python datetime.datetime
407                let timestamp = dt.timestamp();
408                let microseconds = dt.timestamp_subsec_micros();
409
410                // Get UTC timezone from Python's datetime module
411                let dt_module = datetime_module(py).context("Failed to import datetime module")?;
412                let datetime_class = dt_module.getattr("datetime")?;
413                let utc_timezone = dt_module.getattr("timezone")?.getattr("utc")?;
414
415                // Create timezone-aware datetime using fromtimestamp
416                let total_seconds = timestamp as f64 + microseconds as f64 / 1_000_000.0;
417                let py_datetime = datetime_class
418                    .call_method1("fromtimestamp", (total_seconds, utc_timezone))
419                    .context("Failed to create Python datetime from timestamp")?;
420
421                dict.set_item(k, py_datetime)
422                    .context("Could not insert timestamp into python dictionary")?
423            }
424        }
425    }
426
427    Ok(dict)
428}
429
430#[cfg(test)]
431mod tests {
432    use std::{ptr::NonNull, sync::Arc};
433
434    use aligned_vec::{AVec, ConstAlign};
435    use arrow::{
436        array::{
437            ArrayData, ArrayRef, BooleanArray, Float64Array, Int8Array, Int32Array, Int64Array,
438            ListArray, StructArray,
439        },
440        buffer::Buffer,
441    };
442
443    use arrow_schema::{DataType, Field};
444    use dora_node_api::arrow_utils::{decode_arrow_ipc_zero_copy, encode_arrow_ipc};
445    use eyre::{Context, Result};
446
447    fn assert_roundtrip(arrow_array: &ArrayData) -> Result<()> {
448        // The data plane is Arrow-IPC-only: encode to a self-describing IPC
449        // stream, then decode it back from a 128-byte-aligned buffer (as the
450        // receive path does).
451        let ipc_bytes = encode_arrow_ipc(arrow_array)?;
452        let mut sample: AVec<u8, ConstAlign<128>> = AVec::from_slice(128, &ipc_bytes);
453
454        let serialized_deserialized_arrow_array = {
455            let ptr = NonNull::new(sample.as_mut_ptr() as *mut _).unwrap();
456            let len = sample.len();
457
458            let raw_buffer = unsafe {
459                arrow::buffer::Buffer::from_custom_allocation(ptr, len, Arc::new(sample))
460            };
461            decode_arrow_ipc_zero_copy(raw_buffer)?
462        };
463
464        assert_eq!(arrow_array, &serialized_deserialized_arrow_array);
465
466        Ok(())
467    }
468
469    #[test]
470    fn serialize_deserialize_arrow() -> Result<()> {
471        // Int8
472        let arrow_array = Int8Array::from(vec![1, -2, 3, 4]).into();
473        assert_roundtrip(&arrow_array).context("Int8Array roundtrip failed")?;
474
475        // Int64
476        let arrow_array = Int64Array::from(vec![1, -2, 3, 4]).into();
477        assert_roundtrip(&arrow_array).context("Int64Array roundtrip failed")?;
478
479        // Float64
480        let arrow_array = Float64Array::from(vec![1., -2., 3., 4.]).into();
481        assert_roundtrip(&arrow_array).context("Float64Array roundtrip failed")?;
482
483        // Struct
484        let boolean = Arc::new(BooleanArray::from(vec![false, false, true, true]));
485        let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
486
487        let struct_array = StructArray::from(vec![
488            (
489                Arc::new(Field::new("b", DataType::Boolean, false)),
490                boolean as ArrayRef,
491            ),
492            (
493                Arc::new(Field::new("c", DataType::Int32, false)),
494                int as ArrayRef,
495            ),
496        ])
497        .into();
498        assert_roundtrip(&struct_array).context("StructArray roundtrip failed")?;
499
500        // List
501        let value_data = ArrayData::builder(DataType::Int32)
502            .len(8)
503            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7]))
504            .build()
505            .unwrap();
506
507        // Construct a buffer for value offsets, for the nested array:
508        //  [[0, 1, 2], [3, 4, 5], [6, 7]]
509        let value_offsets = Buffer::from_slice_ref([0, 3, 6, 8]);
510
511        // Construct a list array from the above two
512        let list_data_type = DataType::List(Arc::new(Field::new("item", DataType::Int32, false)));
513        let list_data = ArrayData::builder(list_data_type)
514            .len(3)
515            .add_buffer(value_offsets)
516            .add_child_data(value_data)
517            .build()
518            .unwrap();
519        let list_array = ListArray::from(list_data).into();
520        assert_roundtrip(&list_array).context("ListArray roundtrip failed")?;
521
522        Ok(())
523    }
524
525    // Regression tests for dora-rs/adora#147: Python event conversion used
526    // to return "UNKNOWN" with no id for five event types, making fault
527    // tolerance and runtime parameters unusable from Python.
528    mod py_event_types {
529        use super::super::PyEvent;
530        use dora_node_api::{
531            Event,
532            dora_core::config::{DataId, NodeId, OperatorId},
533        };
534
535        #[test]
536        fn node_restarted_has_type_and_id() {
537            let event = Event::NodeRestarted {
538                id: NodeId::from("upstream".to_string()),
539            };
540            assert_eq!(PyEvent::ty(&event), "NODE_RESTARTED");
541            assert_eq!(PyEvent::id(&event), Some("upstream"));
542        }
543
544        #[test]
545        fn input_recovered_has_type_and_id() {
546            let event = Event::InputRecovered {
547                id: DataId::from("sensor".to_string()),
548            };
549            assert_eq!(PyEvent::ty(&event), "INPUT_RECOVERED");
550            assert_eq!(PyEvent::id(&event), Some("sensor"));
551        }
552
553        #[test]
554        fn reload_has_type_and_operator_id() {
555            let event = Event::Reload {
556                operator_id: Some(OperatorId::from("op-1".to_string())),
557            };
558            assert_eq!(PyEvent::ty(&event), "RELOAD");
559            assert_eq!(PyEvent::id(&event), Some("op-1"));
560        }
561
562        #[test]
563        fn reload_without_operator_has_no_id() {
564            let event = Event::Reload { operator_id: None };
565            assert_eq!(PyEvent::ty(&event), "RELOAD");
566            assert_eq!(PyEvent::id(&event), None);
567        }
568
569        #[test]
570        fn param_update_has_type_and_key() {
571            let event = Event::ParamUpdate {
572                key: "threshold".to_string(),
573                value: serde_json::json!(0.85),
574            };
575            assert_eq!(PyEvent::ty(&event), "PARAM_UPDATE");
576            assert_eq!(PyEvent::id(&event), Some("threshold"));
577        }
578
579        #[test]
580        fn param_deleted_has_type_and_key() {
581            let event = Event::ParamDeleted {
582                key: "threshold".to_string(),
583            };
584            assert_eq!(PyEvent::ty(&event), "PARAM_DELETED");
585            assert_eq!(PyEvent::id(&event), Some("threshold"));
586        }
587    }
588}
dora_operator_api_python/lib.rs

dora_operator_api_python/
lib.rs
