Skip to main content

nautilus_persistence/python/backend/
writer.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Generic backend-selected streaming writer, mirroring the catalog backend dispatch.
17
18use 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
41/// Source clock plus the shared atomic the writer reads time from.
42type ClockBridge = (Rc<RefCell<dyn Clock>>, Arc<AtomicU64>);
43
44/// Python binding for the backend-selected streaming writer.
45///
46/// Appends Feather files under a local directory. With a `catalog`, the writer for that catalog's
47/// backend (`Parquet` or a registered name) also promotes the files into it; without one, the
48/// `Feather` writer keeps only the Feather files.
49#[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    /// Present when constructed with a non-live clock: the source clock plus the
59    /// shared atomic the core writer reads, refreshed before each forwarded call.
60    clock_bridge: Option<ClockBridge>,
61}
62
63#[pymethods]
64#[pyo3_stub_gen::derive::gen_stub_pymethods]
65impl PyStreamingWriter {
66    /// Creates a streaming writer for the local `path`, promoting into `catalog` when given.
67    #[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    /// Returns the resolved writer backend name.
102    #[getter]
103    #[must_use]
104    pub fn backend(&self) -> &str {
105        &self.backend
106    }
107
108    /// Writes a single Nautilus data value.
109    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    /// Flushes buffered data to durable storage.
119    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    /// Closes the writer after flushing buffered data.
128    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    // Pushes the source clock's current time into the shared atomic so test clocks
193    // advanced from Python are observed by the core writer.
194    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}