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

nautilus_persistence/backend/
session.rs
