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    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
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/// Resolves the writer through the persistence writer-factory registry, mirroring
47/// `writer_backend` selection: `Feather`, `Parquet`, or a registered name.
48#[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    /// Present when constructed with a non-live clock: the source clock plus the
58    /// shared atomic the core writer reads, refreshed before each forwarded call.
59    clock_bridge: Option<ClockBridge>,
60}
61
62#[pymethods]
63#[pyo3_stub_gen::derive::gen_stub_pymethods]
64impl PyStreamingWriter {
65    /// Creates a streaming writer for the given backend name.
66    #[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    /// Returns the resolved writer backend name.
100    #[getter]
101    #[must_use]
102    pub fn backend(&self) -> &str {
103        &self.backend
104    }
105
106    /// Writes a single Nautilus data value.
107    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    /// Flushes buffered data to durable storage.
117    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    /// Closes the writer after flushing buffered data.
126    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    // Pushes the source clock's current time into the shared atomic so test clocks
185    // advanced from Python are observed by the core writer.
186    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}