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