nautilus_persistence/python/backend/
writer.rs1use std::{
19 cell::RefCell,
20 collections::HashMap,
21 rc::Rc,
22 sync::{Arc, atomic::AtomicU64},
23};
24
25use nautilus_common::{clock::Clock, python::clock::PyClock};
26use nautilus_model::data::{
27 Bar, CustomData, Data, IndexPriceUpdate, InstrumentClose, MarkPriceUpdate, OrderBookDelta,
28 OrderBookDepth, QuoteTick, TradeTick,
29};
30use pyo3::{exceptions::PyIOError, prelude::*};
31
32use crate::{
33 backend::default_writer_factories,
34 writer::{
35 factory::{WriterBackendType, WriterConnectConfig, create_writer},
36 feather::WriterClock,
37 traits::StreamingDataSink,
38 },
39};
40
41type ClockBridge = (Rc<RefCell<dyn Clock>>, Arc<AtomicU64>);
43
44#[pyclass(
49 name = "StreamingWriter",
50 module = "nautilus_trader.persistence",
51 unsendable
52)]
53#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")]
54pub struct PyStreamingWriter {
55 sink: Rc<RefCell<StreamingDataSink>>,
56 backend: String,
57 clock_bridge: Option<ClockBridge>,
60}
61
62#[pymethods]
63#[pyo3_stub_gen::derive::gen_stub_pymethods]
64impl PyStreamingWriter {
65 #[new]
67 #[pyo3(signature = (backend, path, clock, storage_options=None))]
68 #[expect(clippy::needless_pass_by_value)]
69 pub fn py_new(
70 backend: &str,
71 path: String,
72 clock: PyClock,
73 storage_options: Option<HashMap<String, String>>,
74 ) -> PyResult<Self> {
75 let backend_type = backend
76 .parse::<WriterBackendType>()
77 .map_err(|e| PyIOError::new_err(format!("Invalid writer backend: {e}")))?;
78 let clock_rc = clock.clock_rc();
79 let (writer_clock, shared_time) = WriterClock::from_shared_clock(&clock_rc);
80 let config = WriterConnectConfig::new(
81 path,
82 storage_options.map(|options| options.into_iter().collect()),
83 );
84 let sink = create_writer(
85 &backend_type,
86 &config,
87 writer_clock,
88 &default_writer_factories(),
89 )
90 .map_err(|e| PyIOError::new_err(format!("Failed to create writer: {e}")))?;
91
92 Ok(Self {
93 sink: Rc::new(RefCell::new(sink)),
94 backend: backend_type.to_string(),
95 clock_bridge: shared_time.map(|shared| (clock_rc, shared)),
96 })
97 }
98
99 #[getter]
101 #[must_use]
102 pub fn backend(&self) -> &str {
103 &self.backend
104 }
105
106 pub fn write(&self, py: Python, data: Py<PyAny>) -> PyResult<()> {
108 let data = pyobject_to_data(py, data)?;
109 self.refresh_writer_clock();
110 self.sink
111 .borrow_mut()
112 .write_data(data)
113 .map_err(|e| PyIOError::new_err(format!("Failed to write data: {e}")))
114 }
115
116 pub fn flush(&self) -> PyResult<()> {
118 self.refresh_writer_clock();
119 self.sink
120 .borrow_mut()
121 .flush()
122 .map_err(|e| PyIOError::new_err(format!("Failed to flush writer: {e}")))
123 }
124
125 pub fn close(&self) -> PyResult<()> {
127 self.refresh_writer_clock();
128 self.sink
129 .borrow_mut()
130 .close()
131 .map_err(|e| PyIOError::new_err(format!("Failed to close writer: {e}")))
132 }
133}
134
135#[expect(
136 clippy::needless_pass_by_value,
137 reason = "PyO3 transfers ownership of the Python object into this conversion boundary"
138)]
139pub(crate) fn pyobject_to_data(py: Python, data: Py<PyAny>) -> PyResult<Data> {
140 let data = data.bind(py);
141
142 if data.is_instance_of::<QuoteTick>() {
143 return Ok(Data::Quote(data.extract::<QuoteTick>()?));
144 }
145
146 if data.is_instance_of::<TradeTick>() {
147 return Ok(Data::Trade(data.extract::<TradeTick>()?));
148 }
149
150 if data.is_instance_of::<Bar>() {
151 return Ok(Data::Bar(data.extract::<Bar>()?));
152 }
153
154 if data.is_instance_of::<OrderBookDelta>() {
155 return Ok(Data::BookDelta(data.extract::<OrderBookDelta>()?));
156 }
157
158 if data.is_instance_of::<OrderBookDepth>() {
159 return Ok(Data::BookDepth(Box::new(data.extract::<OrderBookDepth>()?)));
160 }
161
162 if data.is_instance_of::<IndexPriceUpdate>() {
163 return Ok(Data::IndexPrice(data.extract::<IndexPriceUpdate>()?));
164 }
165
166 if data.is_instance_of::<MarkPriceUpdate>() {
167 return Ok(Data::MarkPrice(data.extract::<MarkPriceUpdate>()?));
168 }
169
170 if data.is_instance_of::<InstrumentClose>() {
171 return Ok(Data::InstrumentClose(data.extract::<InstrumentClose>()?));
172 }
173
174 if data.is_instance_of::<CustomData>() {
175 return Ok(Data::Custom(data.extract::<CustomData>()?));
176 }
177
178 Err(PyIOError::new_err(
179 "Unsupported data type for streaming writer",
180 ))
181}
182
183impl PyStreamingWriter {
184 fn refresh_writer_clock(&self) {
187 if let Some((clock, shared)) = &self.clock_bridge {
188 shared.store(
189 clock.borrow().timestamp_ns().as_u64(),
190 std::sync::atomic::Ordering::Relaxed,
191 );
192 }
193 }
194}