Skip to main content

shape_runtime/data/
query.rs

1//! Data query specification
2//!
3//! Describes what data to load from a DataProvider.
4
5use super::Timeframe;
6
7/// Query specification for loading data
8#[derive(Debug, Clone)]
9pub struct DataQuery {
10    /// ID to query
11    pub id: String,
12    /// Timeframe of the data
13    pub timeframe: Timeframe,
14    /// Start timestamp (Unix seconds), inclusive
15    pub start: Option<i64>,
16    /// End timestamp (Unix seconds), inclusive
17    pub end: Option<i64>,
18    /// Maximum number of rows to return
19    pub limit: Option<usize>,
20}
21
22impl DataQuery {
23    /// Create a new query for an ID and timeframe
24    pub fn new(id: &str, timeframe: Timeframe) -> Self {
25        Self {
26            id: id.to_string(),
27            timeframe,
28            start: None,
29            end: None,
30            limit: None,
31        }
32    }
33
34    /// Set the time range (builder pattern)
35    pub fn range(mut self, start: i64, end: i64) -> Self {
36        self.start = Some(start);
37        self.end = Some(end);
38        self
39    }
40
41    /// Set the start time
42    pub fn start(mut self, start: i64) -> Self {
43        self.start = Some(start);
44        self
45    }
46
47    /// Set the end time
48    pub fn end(mut self, end: i64) -> Self {
49        self.end = Some(end);
50        self
51    }
52
53    /// Set the limit (builder pattern)
54    pub fn limit(mut self, n: usize) -> Self {
55        self.limit = Some(n);
56        self
57    }
58
59    /// Check if this query has a time range
60    pub fn has_range(&self) -> bool {
61        self.start.is_some() || self.end.is_some()
62    }
63}
64
65impl Default for DataQuery {
66    fn default() -> Self {
67        Self::new("", Timeframe::default())
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn test_query_builder() {
77        let query = DataQuery::new("AAPL", Timeframe::d1())
78            .range(1000, 2000)
79            .limit(100);
80
81        assert_eq!(query.id, "AAPL");
82        assert_eq!(query.start, Some(1000));
83        assert_eq!(query.end, Some(2000));
84        assert_eq!(query.limit, Some(100));
85    }
86}