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

nautilus_persistence/backend/
session.rs
