nautilus_persistence/backend/
session.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
16use std::{
17    sync::{
18        Arc,
19        atomic::{AtomicBool, Ordering},
20    },
21    vec::IntoIter,
22};
23
24use ahash::{AHashMap, AHashSet};
25use datafusion::{
26    arrow::record_batch::RecordBatch,
27    error::{DataFusionError, Result},
28    logical_expr::expr::Sort,
29    physical_plan::SendableRecordBatchStream,
30    prelude::*,
31};
32use futures::{Stream, StreamExt};
33use nautilus_common::live::get_runtime;
34use nautilus_core::UnixNanos;
35use nautilus_model::data::{Data, HasTsInit};
36use nautilus_serialization::arrow::{
37    DataStreamingError, DecodeDataFromRecordBatch, EncodeToRecordBatch, EncodingError, WriteStream,
38};
39use object_store::ObjectStore;
40use parking_lot::Mutex;
41use url::Url;
42
43use super::{
44    compare::Compare,
45    kmerge_batch::{EagerStream, ElementBatchIter, KMerge},
46};
47
48#[derive(Debug, Default)]
49pub struct TsInitComparator;
50
51impl<I> Compare<ElementBatchIter<I, Data>> for TsInitComparator
52where
53    I: Iterator<Item = IntoIter<Data>>,
54{
55    fn compare(
56        &self,
57        l: &ElementBatchIter<I, Data>,
58        r: &ElementBatchIter<I, Data>,
59    ) -> std::cmp::Ordering {
60        // Max heap ordering must be reversed
61        l.item.ts_init().cmp(&r.item.ts_init()).reverse()
62    }
63}
64
65/// Represents a failure raised by a query's underlying data stream.
66#[derive(Debug, thiserror::Error)]
67pub enum QueryError {
68    /// The record batch stream returned an error.
69    #[error("Record batch stream error: {0}")]
70    Stream(#[from] DataFusionError),
71    /// A record batch could not be decoded into Nautilus data.
72    #[error("Record batch decode error: {0}")]
73    Decode(#[from] EncodingError),
74}
75
76/// Holds the first failure observed by any of a query's batch streams.
77///
78/// `failed` keeps the common path off the mutex, because [`QueryResult::next`] consults the slot
79/// once per merged item while loading a catalog.
80#[derive(Default)]
81struct ErrorSlot {
82    failed: AtomicBool,
83    error: Mutex<Option<QueryError>>,
84}
85
86impl ErrorSlot {
87    fn record(&self, error: QueryError) {
88        self.error.lock().get_or_insert(error);
89        self.failed.store(true, Ordering::Release);
90    }
91
92    fn failed(&self) -> bool {
93        self.failed.load(Ordering::Acquire)
94    }
95
96    fn take(&self) -> Option<QueryError> {
97        self.error.lock().take()
98    }
99}
100
101/// Iterates the merged data of every registered query stream in ascending `ts_init` order.
102///
103/// A batch stream that fails stops contributing data and the failure is yielded as an error, so a
104/// failed query can never be mistaken for an exhausted one.
105pub struct QueryResult {
106    merge: KMerge<BatchStream, Data, TsInitComparator>,
107    error: Arc<ErrorSlot>,
108}
109
110impl QueryResult {
111    /// Discards the remaining data streams without draining them.
112    pub fn clear(&mut self) {
113        self.merge.clear();
114    }
115}
116
117impl Iterator for QueryResult {
118    // Spelled out because `Result` is the DataFusion alias in this module
119    type Item = std::result::Result<Data, QueryError>;
120
121    fn next(&mut self) -> Option<Self::Item> {
122        // A failure recorded while merging the previous item ends the query, so drop the
123        // remaining streams rather than returning data from an incomplete result.
124        if self.error.failed()
125            && let Some(e) = self.error.take()
126        {
127            self.clear();
128            return Some(Err(e));
129        }
130
131        match self.merge.next() {
132            Some(item) => Some(Ok(item)),
133            // Always taken, so a failure recorded on the final poll cannot read as exhaustion
134            None => self.error.take().map(Err),
135        }
136    }
137}
138
139/// Provides a DataFusion session and registers DataFusion queries.
140///
141/// The session is used to register data sources and make queries on them. A
142/// query returns a Chunk of Arrow records. It is decoded and converted into
143/// a Vec of data by types that implement [`DecodeDataFromRecordBatch`].
144#[cfg_attr(
145    feature = "python",
146    pyo3::pyclass(module = "nautilus_trader.persistence", unsendable)
147)]
148#[cfg_attr(
149    feature = "python",
150    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")
151)]
152pub struct DataBackendSession {
153    pub chunk_size: usize,
154    pub runtime: tokio::runtime::Handle,
155    session_ctx: SessionContext,
156    batch_streams: Vec<BatchStream>,
157    error: Arc<ErrorSlot>,
158    registered_tables: AHashSet<String>,
159}
160
161impl DataBackendSession {
162    /// Creates a new [`DataBackendSession`] instance.
163    #[must_use]
164    pub fn new(chunk_size: usize) -> Self {
165        let session_cfg = SessionConfig::new()
166            .set_str("datafusion.optimizer.repartition_file_scans", "false")
167            .set_str("datafusion.optimizer.prefer_existing_sort", "true");
168        let session_ctx = SessionContext::new_with_config(session_cfg);
169        Self {
170            session_ctx,
171            batch_streams: Vec::default(),
172            error: Arc::default(),
173            chunk_size,
174            runtime: get_runtime().handle().clone(),
175            registered_tables: AHashSet::new(),
176        }
177    }
178
179    /// Register an object store with the session context
180    pub fn register_object_store(&mut self, url: &Url, object_store: Arc<dyn ObjectStore>) {
181        self.session_ctx.register_object_store(url, object_store);
182    }
183
184    /// Register an object store with the session context from a URI with optional storage options.
185    ///
186    /// # Errors
187    ///
188    /// Returns an error if the object store URI cannot be normalized or the backend
189    /// cannot be created.
190    pub fn register_object_store_from_uri(
191        &mut self,
192        uri: &str,
193        storage_options: Option<AHashMap<String, String>>,
194    ) -> anyhow::Result<()> {
195        let location =
196            crate::parquet::create_object_store_location_from_path(uri, storage_options)?;
197
198        if let Some(root_url) = location.store_root_url().cloned() {
199            self.register_object_store(&root_url, location.object_store);
200        }
201
202        Ok(())
203    }
204
205    /// Writes encoded data to a streaming sink.
206    ///
207    /// # Errors
208    ///
209    /// Returns an error if Arrow encoding or stream writing fails.
210    pub fn write_data<T: EncodeToRecordBatch>(
211        data: &[T],
212        metadata: &AHashMap<String, String>,
213        stream: &mut dyn WriteStream,
214    ) -> Result<(), DataStreamingError> {
215        // Convert AHashMap to HashMap for Arrow compatibility
216        let metadata: std::collections::HashMap<String, String> = metadata
217            .iter()
218            .map(|(k, v)| (k.clone(), v.clone()))
219            .collect();
220        let record_batch = T::encode_batch(&metadata, data)?;
221        stream.write(&record_batch)?;
222        Ok(())
223    }
224
225    /// Registers a Parquet file and adds a batch stream for decoding.
226    ///
227    /// The caller must specify `T` to indicate the kind of data expected. `table_name` is
228    /// the logical name for queries; `file_path` is the Parquet path; `sql_query` defaults
229    /// to `SELECT * FROM {table_name} ORDER BY ts_init` if `None`.
230    ///
231    /// When `custom_type_name` is `Some`, it is merged into each batch's schema metadata
232    /// before decoding (as `type_name`). Use this for custom data when Parquet/DataFusion
233    /// does not preserve schema metadata so the decoder can look up the type in the registry.
234    ///
235    /// The file data must be ordered by the `ts_init` in ascending order for this
236    /// to work correctly.
237    ///
238    /// # Errors
239    ///
240    /// Returns an error if parquet registration, SQL planning, stream execution, or
241    /// data decoding setup fails.
242    pub fn add_file<T>(
243        &mut self,
244        table_name: &str,
245        file_path: &str,
246        sql_query: Option<&str>,
247        custom_type_name: Option<&str>,
248    ) -> Result<()>
249    where
250        T: DecodeDataFromRecordBatch,
251    {
252        // Check if table is already registered to avoid duplicates
253        let is_new_table = !self.registered_tables.contains(table_name);
254
255        if is_new_table {
256            // Register the table only if it doesn't exist
257            let parquet_options = ParquetReadOptions::<'_> {
258                skip_metadata: Some(false),
259                file_sort_order: vec![vec![Sort {
260                    expr: col("ts_init"),
261                    asc: true,
262                    nulls_first: false,
263                }]],
264                ..Default::default()
265            };
266            super::block_on(
267                &self.runtime,
268                self.session_ctx
269                    .register_parquet(table_name, file_path, parquet_options),
270            )?;
271
272            self.registered_tables.insert(table_name.to_string());
273
274            // Only add batch stream for newly registered tables to avoid duplicates
275            let default_query = format!("SELECT * FROM {table_name} ORDER BY ts_init");
276            let sql_query = sql_query.unwrap_or(&default_query);
277            let query = super::block_on(&self.runtime, self.session_ctx.sql(sql_query))?;
278            let batch_stream = super::block_on(&self.runtime, query.execute_stream())?;
279            self.add_batch_stream::<T>(batch_stream, custom_type_name.map(String::from));
280        }
281
282        Ok(())
283    }
284
285    /// Registers a Parquet file and executes a query, returning the raw record batches.
286    ///
287    /// # Errors
288    ///
289    /// Returns an error if parquet registration, SQL planning, stream execution, or
290    /// batch collection fails.
291    pub fn collect_query_batches(
292        &mut self,
293        table_name: &str,
294        file_path: &str,
295        sql_query: Option<&str>,
296    ) -> Result<Vec<RecordBatch>> {
297        if !self.registered_tables.contains(table_name) {
298            let parquet_options = ParquetReadOptions::<'_> {
299                skip_metadata: Some(false),
300                file_sort_order: vec![vec![Sort {
301                    expr: col("ts_init"),
302                    asc: true,
303                    nulls_first: false,
304                }]],
305                ..Default::default()
306            };
307            super::block_on(
308                &self.runtime,
309                self.session_ctx
310                    .register_parquet(table_name, file_path, parquet_options),
311            )?;
312
313            self.registered_tables.insert(table_name.to_string());
314        }
315
316        let default_query = format!("SELECT * FROM {table_name} ORDER BY ts_init");
317        let sql_query = sql_query.unwrap_or(&default_query);
318        let query = super::block_on(&self.runtime, self.session_ctx.sql(sql_query))?;
319        let mut batch_stream = super::block_on(&self.runtime, query.execute_stream())?;
320
321        super::block_on(&self.runtime, async {
322            let mut batches = Vec::new();
323            while let Some(batch) = batch_stream.next().await {
324                batches.push(batch?);
325            }
326            Ok::<_, datafusion::error::DataFusionError>(batches)
327        })
328    }
329
330    fn add_batch_stream<T>(
331        &mut self,
332        stream: SendableRecordBatchStream,
333        custom_type_name: Option<String>,
334    ) where
335        T: DecodeDataFromRecordBatch,
336    {
337        self.batch_streams.push(BatchStream {
338            inner: EagerStream::from_stream_with_runtime(
339                decode_batches::<T>(stream, custom_type_name),
340                self.runtime.clone(),
341            ),
342            error: Arc::clone(&self.error),
343        });
344    }
345
346    // Consumes the registered queries and returns a [`QueryResult].
347    // Passes the output of the query though the a KMerge which sorts the
348    // queries in ascending order of `ts_init`.
349    // QueryResult is an iterator that return Vec<Data>.
350    pub fn get_query_result(&mut self) -> QueryResult {
351        let mut merge: KMerge<_, _, _> = KMerge::new(TsInitComparator);
352
353        self.batch_streams
354            .drain(..)
355            .for_each(|batch_stream| merge.push_iter(batch_stream));
356
357        QueryResult {
358            merge,
359            error: std::mem::take(&mut self.error),
360        }
361    }
362
363    /// Clears all registered tables and batch streams.
364    ///
365    /// This is useful when the underlying files have changed and we need to
366    /// re-register tables with updated data.
367    pub fn clear_registered_tables(&mut self) {
368        self.registered_tables.clear();
369        self.batch_streams.clear();
370        self.error = Arc::default();
371
372        // Create a new session context to completely reset the DataFusion state
373        let session_cfg = SessionConfig::new()
374            .set_str("datafusion.optimizer.repartition_file_scans", "false")
375            .set_str("datafusion.optimizer.prefer_existing_sort", "true");
376        self.session_ctx = SessionContext::new_with_config(session_cfg);
377    }
378}
379
380type BatchResult = std::result::Result<IntoIter<Data>, QueryError>;
381
382/// Decodes each record batch, yielding the first failure and then ending the stream.
383///
384/// A record batch stream that has returned an error gives no guarantee about being polled again,
385/// and a panic in the producer task would abort the process under `panic = "abort"`.
386fn decode_batches<T>(
387    stream: SendableRecordBatchStream,
388    custom_type_name: Option<String>,
389) -> impl Stream<Item = BatchResult> + Send + 'static
390where
391    T: DecodeDataFromRecordBatch,
392{
393    futures::stream::unfold(
394        (stream, custom_type_name, false),
395        |(mut stream, custom_type_name, failed)| async move {
396            if failed {
397                return None;
398            }
399
400            let batch = decode_batch::<T>(stream.next().await?, custom_type_name.as_deref());
401            let failed = batch.is_err();
402
403            Some((batch, (stream, custom_type_name, failed)))
404        },
405    )
406}
407
408fn decode_batch<T>(
409    result: std::result::Result<RecordBatch, DataFusionError>,
410    custom_type_name: Option<&str>,
411) -> BatchResult
412where
413    T: DecodeDataFromRecordBatch,
414{
415    let batch = result?;
416    let mut metadata: std::collections::HashMap<String, String> = batch.schema().metadata().clone();
417
418    if let Some(type_name) = custom_type_name {
419        metadata.insert("type_name".to_string(), type_name.to_string());
420    }
421
422    Ok(T::decode_data_batch(&metadata, batch)?.into_iter())
423}
424
425/// Feeds decoded batches to the merge and diverts a failure to the shared error slot.
426///
427/// The merge orders items by `ts_init`, so a failure cannot travel with the data. Recording it
428/// here ends this stream for the merge while leaving the remaining streams intact, and lets
429/// [`QueryResult`] report the failure instead of exhaustion.
430struct BatchStream {
431    inner: EagerStream<BatchResult>,
432    error: Arc<ErrorSlot>,
433}
434
435impl Iterator for BatchStream {
436    type Item = IntoIter<Data>;
437
438    fn next(&mut self) -> Option<Self::Item> {
439        match self.inner.next()? {
440            Ok(batch) => Some(batch),
441            Err(e) => {
442                self.error.record(e);
443                None
444            }
445        }
446    }
447}
448
449#[must_use]
450pub fn build_query(
451    table: &str,
452    start: Option<UnixNanos>,
453    end: Option<UnixNanos>,
454    where_clause: Option<&str>,
455) -> String {
456    let mut conditions = Vec::new();
457
458    // Add where clause if provided
459    if let Some(clause) = where_clause {
460        conditions.push(clause.to_string());
461    }
462
463    // Add start condition if provided
464    if let Some(start_ts) = start {
465        conditions.push(format!("ts_init >= {start_ts}"));
466    }
467
468    // Add end condition if provided
469    if let Some(end_ts) = end {
470        conditions.push(format!("ts_init <= {end_ts}"));
471    }
472
473    // Build base query
474    let mut query = format!("SELECT * FROM {table}");
475
476    // Add WHERE clause if there are conditions
477    if !conditions.is_empty() {
478        query.push_str(" WHERE ");
479        query.push_str(&conditions.join(" AND "));
480    }
481
482    // Add ORDER BY clause
483    query.push_str(" ORDER BY ts_init");
484
485    query
486}
487
488#[cfg_attr(
489    feature = "python",
490    pyo3::pyclass(module = "nautilus_trader.persistence", unsendable)
491)]
492#[cfg_attr(
493    feature = "python",
494    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")
495)]
496pub struct DataQueryResult {
497    pub result: QueryResult,
498    pub acc: Vec<Data>,
499    pub size: usize,
500}
501
502impl DataQueryResult {
503    /// Creates a new [`DataQueryResult`] instance.
504    #[must_use]
505    pub const fn new(result: QueryResult, size: usize) -> Self {
506        Self {
507            result,
508            acc: Vec::new(),
509            size,
510        }
511    }
512}
513
514impl Iterator for DataQueryResult {
515    // An empty chunk signals exhaustion, so a failure must be reported as an error
516    type Item = std::result::Result<Vec<Data>, QueryError>;
517
518    fn next(&mut self) -> Option<Self::Item> {
519        // Poll at least once, since a zero chunk size would return an empty chunk without ever
520        // consulting the query, hiding a failure behind the exhaustion signal.
521        let size = self.size.max(1);
522
523        for _ in 0..size {
524            match self.result.next() {
525                Some(Ok(item)) => self.acc.push(item),
526                Some(Err(e)) => {
527                    self.acc.clear();
528                    return Some(Err(e));
529                }
530                None => break,
531            }
532        }
533
534        // TODO: consider using drain here if perf is unchanged
535        // Some(self.acc.drain(0..).collect())
536        let mut acc: Vec<Data> = Vec::new();
537        std::mem::swap(&mut acc, &mut self.acc);
538        Some(Ok(acc))
539    }
540}
541
542impl Drop for DataQueryResult {
543    fn drop(&mut self) {
544        self.result.clear();
545    }
546}
547
548#[cfg(test)]
549mod tests {
550    use std::{collections::HashMap, sync::atomic::AtomicUsize, task::Poll};
551
552    use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
553    use nautilus_common::live::get_runtime;
554    use nautilus_model::{
555        data::QuoteTick,
556        identifiers::InstrumentId,
557        types::{Price, Quantity},
558    };
559    use nautilus_serialization::arrow::{
560        ArrowSchemaProvider, KEY_INSTRUMENT_ID, KEY_PRICE_PRECISION, KEY_SIZE_PRECISION,
561    };
562    #[cfg(feature = "python")]
563    use pyo3::{Py, Python, exceptions::PyRuntimeError, types::PyAnyMethods};
564    use rstest::rstest;
565
566    use super::*;
567
568    const INSTRUMENT_ID: &str = "EUR/USD.SIM";
569
570    fn quote(ts_init: u64) -> QuoteTick {
571        QuoteTick::new(
572            InstrumentId::from(INSTRUMENT_ID),
573            Price::from("1.0001"),
574            Price::from("1.0002"),
575            Quantity::from("100"),
576            Quantity::from("100"),
577            UnixNanos::from(ts_init),
578            UnixNanos::from(ts_init),
579        )
580    }
581
582    fn quote_metadata() -> HashMap<String, String> {
583        HashMap::from([
584            (KEY_INSTRUMENT_ID.to_string(), INSTRUMENT_ID.to_string()),
585            (KEY_PRICE_PRECISION.to_string(), "4".to_string()),
586            (KEY_SIZE_PRECISION.to_string(), "0".to_string()),
587        ])
588    }
589
590    fn quote_batch(quotes: &[QuoteTick]) -> RecordBatch {
591        QuoteTick::encode_batch(&quote_metadata(), quotes).expect("failed to encode quotes")
592    }
593
594    fn stream_error() -> DataFusionError {
595        DataFusionError::Execution("injected stream failure".to_string())
596    }
597
598    fn batch_stream(
599        batches: Vec<std::result::Result<RecordBatch, DataFusionError>>,
600    ) -> SendableRecordBatchStream {
601        Box::pin(RecordBatchStreamAdapter::new(
602            Arc::new(QuoteTick::get_schema(Some(quote_metadata()))),
603            futures::stream::iter(batches),
604        ))
605    }
606
607    fn ts_inits(items: &[std::result::Result<Data, QueryError>]) -> Vec<u64> {
608        items
609            .iter()
610            .filter_map(|item| item.as_ref().ok())
611            .map(|data| data.ts_init().as_u64())
612            .collect()
613    }
614
615    #[rstest]
616    fn data_backend_sessions_share_global_runtime() {
617        let first = DataBackendSession::new(10);
618        let second = DataBackendSession::new(10);
619
620        assert_eq!(first.runtime.id(), second.runtime.id());
621        assert_eq!(first.runtime.id(), get_runtime().handle().id());
622    }
623
624    #[rstest]
625    fn query_result_merges_streams_in_order_then_exhausts() {
626        let mut session = DataBackendSession::new(10);
627        session.add_batch_stream::<QuoteTick>(
628            batch_stream(vec![
629                Ok(quote_batch(&[quote(1), quote(3)])),
630                Ok(quote_batch(&[quote(5)])),
631            ]),
632            None,
633        );
634        session.add_batch_stream::<QuoteTick>(
635            batch_stream(vec![Ok(quote_batch(&[quote(2), quote(4)]))]),
636            None,
637        );
638
639        let mut result = session.get_query_result();
640        let items: Vec<_> = result.by_ref().collect();
641
642        assert_eq!(ts_inits(&items), vec![1, 2, 3, 4, 5]);
643        assert_eq!(items.len(), 5);
644        assert!(result.next().is_none());
645    }
646
647    #[rstest]
648    fn query_result_reports_stream_error_after_its_data() {
649        let mut session = DataBackendSession::new(10);
650        session.add_batch_stream::<QuoteTick>(
651            batch_stream(vec![
652                Ok(quote_batch(&[quote(1), quote(2)])),
653                Err(stream_error()),
654            ]),
655            None,
656        );
657
658        let mut result = session.get_query_result();
659        let items: Vec<_> = result.by_ref().collect();
660
661        assert_eq!(ts_inits(&items), vec![1, 2]);
662        assert_eq!(items.len(), 3);
663        assert!(
664            matches!(items[2], Err(QueryError::Stream(_))),
665            "expected a stream error, was {:?}",
666            items[2]
667        );
668        assert!(result.next().is_none());
669    }
670
671    #[rstest]
672    fn query_result_stops_when_one_of_many_streams_fails() {
673        let mut session = DataBackendSession::new(10);
674        session.add_batch_stream::<QuoteTick>(
675            batch_stream(vec![Ok(quote_batch(&[quote(1)])), Err(stream_error())]),
676            None,
677        );
678        session
679            .add_batch_stream::<QuoteTick>(batch_stream(vec![Ok(quote_batch(&[quote(2)]))]), None);
680
681        let mut result = session.get_query_result();
682        let items: Vec<_> = result.by_ref().collect();
683
684        // The healthy stream still holds `quote(2)`, so exhaustion here would look successful
685        assert_eq!(ts_inits(&items), vec![1]);
686        assert_eq!(items.len(), 2);
687        assert!(
688            matches!(items[1], Err(QueryError::Stream(_))),
689            "expected a stream error, was {:?}",
690            items[1]
691        );
692        assert!(result.next().is_none());
693    }
694
695    #[rstest]
696    fn query_result_reports_decode_error() {
697        let mut session = DataBackendSession::new(10);
698        // Encode without schema metadata so the decoder cannot resolve the instrument
699        let batch =
700            QuoteTick::encode_batch(&HashMap::new(), &[quote(1)]).expect("failed to encode quotes");
701        session.add_batch_stream::<QuoteTick>(batch_stream(vec![Ok(batch)]), None);
702
703        let items: Vec<_> = session.get_query_result().collect();
704
705        assert_eq!(items.len(), 1);
706        assert!(
707            matches!(
708                items[0],
709                Err(QueryError::Decode(EncodingError::MissingMetadata(
710                    KEY_INSTRUMENT_ID
711                )))
712            ),
713            "expected a decode error, was {:?}",
714            items[0]
715        );
716    }
717
718    #[rstest]
719    fn data_query_result_reports_error_instead_of_an_empty_chunk() {
720        let mut session = DataBackendSession::new(10);
721        session.add_batch_stream::<QuoteTick>(
722            batch_stream(vec![Ok(quote_batch(&[quote(1)])), Err(stream_error())]),
723            None,
724        );
725
726        let mut result = DataQueryResult::new(session.get_query_result(), 10);
727        let chunk = result.next().expect("chunked result must yield an item");
728
729        assert!(
730            matches!(chunk, Err(QueryError::Stream(_))),
731            "expected a stream error, was {chunk:?}"
732        );
733
734        let after = result
735            .next()
736            .expect("chunked result must signal exhaustion")
737            .expect("a failed query must not fail twice");
738
739        assert!(after.is_empty(), "the discarded chunk must not be replayed");
740    }
741
742    #[rstest]
743    fn data_query_result_reports_an_error_with_a_zero_chunk_size() {
744        let mut session = DataBackendSession::new(10);
745        session.add_batch_stream::<QuoteTick>(batch_stream(vec![Err(stream_error())]), None);
746
747        let mut result = DataQueryResult::new(session.get_query_result(), 0);
748        let chunk = result.next().expect("chunked result must yield an item");
749
750        assert!(
751            matches!(chunk, Err(QueryError::Stream(_))),
752            "expected a stream error, was {chunk:?}"
753        );
754    }
755
756    #[rstest]
757    fn data_query_result_ends_with_an_empty_chunk_when_successful() {
758        let mut session = DataBackendSession::new(10);
759        session.add_batch_stream::<QuoteTick>(
760            batch_stream(vec![Ok(quote_batch(&[quote(1), quote(2)]))]),
761            None,
762        );
763
764        let mut result = DataQueryResult::new(session.get_query_result(), 10);
765        let chunk = result
766            .next()
767            .expect("chunked result must yield a chunk")
768            .expect("query must not fail");
769
770        assert_eq!(chunk.len(), 2);
771        assert_eq!(
772            chunk
773                .iter()
774                .map(|data| data.ts_init().as_u64())
775                .collect::<Vec<_>>(),
776            vec![1, 2]
777        );
778
779        let last = result
780            .next()
781            .expect("chunked result must signal exhaustion")
782            .expect("query must not fail");
783
784        assert!(last.is_empty());
785    }
786
787    #[rstest]
788    fn decode_batches_stops_polling_a_failed_stream() {
789        let polls = Arc::new(AtomicUsize::new(0));
790        let counted = Arc::clone(&polls);
791        let mut batches = vec![Ok(quote_batch(&[quote(1)])), Err(stream_error())].into_iter();
792        let inner = futures::stream::poll_fn(move |_| {
793            counted.fetch_add(1, Ordering::SeqCst);
794            Poll::Ready(batches.next())
795        });
796        let stream = Box::pin(RecordBatchStreamAdapter::new(
797            Arc::new(QuoteTick::get_schema(Some(quote_metadata()))),
798            inner,
799        ));
800
801        let decoded = decode_batches::<QuoteTick>(stream, None);
802        let items: Vec<_> = futures::executor::block_on_stream(Box::pin(decoded)).collect();
803
804        assert_eq!(items.len(), 2);
805        assert!(items[0].is_ok());
806        assert!(matches!(items[1], Err(QueryError::Stream(_))));
807        assert_eq!(polls.load(Ordering::SeqCst), 2);
808    }
809
810    #[rstest]
811    fn a_new_query_does_not_inherit_an_earlier_failure() {
812        let mut session = DataBackendSession::new(10);
813        session.add_batch_stream::<QuoteTick>(batch_stream(vec![Err(stream_error())]), None);
814
815        // Held open so a shared error slot would leak into the query registered next
816        let failed = session.get_query_result();
817
818        session.add_batch_stream::<QuoteTick>(
819            batch_stream(vec![Ok(quote_batch(&[quote(1), quote(2)]))]),
820            None,
821        );
822        let items: Vec<_> = session.get_query_result().collect();
823
824        let failed: Vec<_> = failed.collect();
825
826        assert_eq!(ts_inits(&items), vec![1, 2]);
827        assert_eq!(items.len(), 2);
828        assert_eq!(failed.len(), 1);
829        assert!(
830            matches!(failed[0], Err(QueryError::Stream(_))),
831            "expected a stream error, was {:?}",
832            failed[0]
833        );
834    }
835
836    #[rstest]
837    #[cfg(feature = "python")]
838    fn python_to_list_raises_on_stream_error() {
839        let mut session = DataBackendSession::new(10);
840        session.add_batch_stream::<QuoteTick>(
841            batch_stream(vec![Ok(quote_batch(&[quote(1)])), Err(stream_error())]),
842            None,
843        );
844        let result = DataQueryResult::new(session.get_query_result(), 10);
845
846        Python::initialize();
847        Python::attach(|py| {
848            let result = Py::new(py, result).expect("failed to create the query result");
849            let error = result
850                .bind(py)
851                .call_method0("to_list")
852                .expect_err("to_list must raise when a stream fails");
853
854            assert!(error.is_instance_of::<PyRuntimeError>(py));
855            assert!(
856                error.to_string().contains("Record batch stream error"),
857                "was {error}"
858            );
859        });
860    }
861
862    #[rstest]
863    #[cfg(feature = "python")]
864    fn python_next_raises_on_decode_error() {
865        let mut session = DataBackendSession::new(10);
866        // Encode without schema metadata so the decoder cannot resolve the instrument
867        let batch =
868            QuoteTick::encode_batch(&HashMap::new(), &[quote(1)]).expect("failed to encode quotes");
869        session.add_batch_stream::<QuoteTick>(batch_stream(vec![Ok(batch)]), None);
870        let result = DataQueryResult::new(session.get_query_result(), 10);
871
872        Python::initialize();
873        Python::attach(|py| {
874            let result = Py::new(py, result).expect("failed to create the query result");
875            let error = result
876                .bind(py)
877                .call_method0("__next__")
878                .expect_err("__next__ must raise when a batch cannot be decoded");
879
880            assert!(error.is_instance_of::<PyRuntimeError>(py));
881            assert!(
882                error.to_string().contains("Record batch decode error"),
883                "was {error}"
884            );
885        });
886    }
887}
nautilus_persistence/backend/session.rs

nautilus_persistence/backend/
session.rs
