query/
options.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
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
15use std::collections::HashMap;
16
17use chrono::{DateTime, Utc};
18use common_base::memory_limit::MemoryLimit;
19use datafusion::config::{ConfigEntry, ConfigExtension, ExtensionOptions};
20use serde::{Deserialize, Serialize};
21use session::context::QueryContextRef;
22use store_api::storage::RegionId;
23use table::metadata::TableId;
24
25use crate::error::{Error, InvalidQueryContextExtensionSnafu, Result};
26
27pub const FLOW_INCREMENTAL_AFTER_SEQS: &str = "flow.incremental_after_seqs";
28pub const FLOW_INCREMENTAL_MODE: &str = "flow.incremental_mode";
29pub const FLOW_RETURN_REGION_SEQ: &str = "flow.return_region_seq";
30pub const FLOW_SINK_TABLE_ID: &str = "flow.sink_table_id";
31/// Flow scheduler binding for the logical time of one scheduled attempt.
32/// Query planning, SQL/TQL parsing, range-select rewrite and DataFusion
33/// execution read this extension so `now()` is stable for the whole attempt.
34pub const FLOW_SCHEDULED_TIME_MILLIS: &str = "flow.scheduled_time_millis";
35/// Enable by default, set to false to explicitly disable.
36pub const QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN: &str =
37    "query.enable_remote_dynamic_filter_pushdown";
38
39pub const FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY: &str = "memtable_only";
40
41/// Query engine config
42#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
43#[serde(default)]
44pub struct QueryOptions {
45    /// Parallelism of query engine. Default to 0, which implies the number of logical CPUs.
46    pub parallelism: usize,
47    /// Whether to allow query fallback when push down fails.
48    pub allow_query_fallback: bool,
49    /// Memory pool size for query execution. Setting it to 0 disables the limit (unbounded).
50    /// Supports absolute size (e.g., "2GB") or percentage (e.g., "50%").
51    /// When this limit is reached, queries will fail with ResourceExhausted error.
52    pub memory_pool_size: MemoryLimit,
53    /// Whether to expose per-region query load metrics.
54    #[serde(skip)]
55    pub enable_per_region_metrics: bool,
56}
57
58#[allow(clippy::derivable_impls)]
59impl Default for QueryOptions {
60    fn default() -> Self {
61        Self {
62            parallelism: 0,
63            allow_query_fallback: false,
64            memory_pool_size: MemoryLimit::default(),
65            enable_per_region_metrics: false,
66        }
67    }
68}
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub enum FlowIncrementalMode {
72    MemtableOnly,
73}
74
75#[derive(Debug, Clone, PartialEq, Eq, Default)]
76pub struct FlowQueryExtensions {
77    /// Maps region id -> lower exclusive sequence bound for incremental reads.
78    pub incremental_after_seqs: Option<HashMap<u64, u64>>,
79    /// Incremental read mode requested by the caller.
80    pub incremental_mode: Option<FlowIncrementalMode>,
81    /// Whether the caller expects per-region watermark metadata in terminal metrics.
82    pub return_region_seq: bool,
83    /// Optional sink table id used to distinguish source scans from sink reads.
84    pub sink_table_id: Option<TableId>,
85}
86
87impl FlowQueryExtensions {
88    /// Parses flow-specific query extensions when any flow key is present.
89    ///
90    /// Returns `Ok(None)` for ordinary queries with no flow-related extensions,
91    /// `Ok(Some(_))` when flow context is present and valid, and `Err(_)` when a
92    /// flow-related extension is present but malformed or incomplete.
93    pub fn parse_flow_extensions(extensions: &HashMap<String, String>) -> Result<Option<Self>> {
94        let has_flow_context = extensions.contains_key(FLOW_INCREMENTAL_AFTER_SEQS)
95            || extensions.contains_key(FLOW_INCREMENTAL_MODE)
96            || extensions.contains_key(FLOW_RETURN_REGION_SEQ)
97            || extensions.contains_key(FLOW_SINK_TABLE_ID);
98
99        if !has_flow_context {
100            return Ok(None);
101        }
102
103        let incremental_mode = extensions
104            .get(FLOW_INCREMENTAL_MODE)
105            .map(|value| match value.as_str() {
106                v if v.eq_ignore_ascii_case(FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY) => {
107                    Ok(FlowIncrementalMode::MemtableOnly)
108                }
109                _ => Err(invalid_query_context_extension(format!(
110                    "Invalid value for {}: {}",
111                    FLOW_INCREMENTAL_MODE, value
112                ))),
113            })
114            .transpose()?;
115
116        let incremental_after_seqs = extensions
117            .get(FLOW_INCREMENTAL_AFTER_SEQS)
118            .map(|value| parse_incremental_after_seqs(value.as_str()))
119            .transpose()?;
120
121        let return_region_seq = extensions
122            .get(FLOW_RETURN_REGION_SEQ)
123            .map(|value| parse_bool(FLOW_RETURN_REGION_SEQ, value.as_str()))
124            .transpose()?
125            .unwrap_or(false);
126
127        let sink_table_id = extensions
128            .get(FLOW_SINK_TABLE_ID)
129            .map(|value| {
130                value.parse::<TableId>().map_err(|_| {
131                    invalid_query_context_extension(format!(
132                        "Invalid value for {}: {}",
133                        FLOW_SINK_TABLE_ID, value
134                    ))
135                })
136            })
137            .transpose()?;
138
139        if matches!(incremental_mode, Some(FlowIncrementalMode::MemtableOnly)) {
140            let after_seqs = incremental_after_seqs.as_ref().ok_or_else(|| {
141                invalid_query_context_extension(format!(
142                    "{} is required when {}={}.",
143                    FLOW_INCREMENTAL_AFTER_SEQS,
144                    FLOW_INCREMENTAL_MODE,
145                    FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY
146                ))
147            })?;
148            if after_seqs.is_empty() {
149                return Err(invalid_query_context_extension(format!(
150                    "{} must not be empty when {}={}.",
151                    FLOW_INCREMENTAL_AFTER_SEQS,
152                    FLOW_INCREMENTAL_MODE,
153                    FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY
154                )));
155            }
156        }
157
158        Ok(Some(Self {
159            incremental_after_seqs,
160            incremental_mode,
161            return_region_seq,
162            sink_table_id,
163        }))
164    }
165
166    pub fn validate_for_scan(&self, source_region_id: RegionId) -> Result<bool> {
167        if self.sink_table_id.is_some() && self.sink_table_id == Some(source_region_id.table_id()) {
168            return Ok(false);
169        }
170
171        if matches!(
172            self.incremental_mode,
173            Some(FlowIncrementalMode::MemtableOnly)
174        ) {
175            let after_seqs = self.incremental_after_seqs.as_ref().ok_or_else(|| {
176                invalid_query_context_extension(format!(
177                    "{} is required when {}=memtable_only.",
178                    FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE
179                ))
180            })?;
181
182            if !after_seqs.contains_key(&source_region_id.as_u64()) {
183                return Err(invalid_query_context_extension(format!(
184                    "Missing region {} in {} when {}=memtable_only.",
185                    source_region_id, FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE
186                )));
187            }
188        }
189
190        Ok(self.incremental_after_seqs.is_some())
191    }
192
193    pub fn should_collect_region_watermark(&self) -> bool {
194        should_collect_region_watermark(
195            self.return_region_seq,
196            self.incremental_after_seqs.is_some(),
197        )
198    }
199}
200
201/// Returns whether query-level remote dynamic filter propagation is enabled.
202///
203/// The option defaults to enabled to preserve existing behavior. Callers may set
204/// `query.enable_remote_dynamic_filter_pushdown=false` in query context
205/// extensions to disable FE->DN remote dynamic filter propagation for a single
206/// query.
207pub fn remote_dyn_filter_pushdown_enabled_from_extensions(
208    extensions: &HashMap<String, String>,
209) -> Result<bool> {
210    extensions
211        .get(QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN)
212        .map(|value| parse_bool(QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN, value.as_str()))
213        .transpose()
214        .map(|value| value.unwrap_or(true))
215}
216
217/// Returns whether raw Flow query extensions request terminal region watermark collection.
218///
219/// This is only an intent/presence check for transport/scan plumbing; callers that need
220/// validated Flow options must still use [`FlowQueryExtensions::parse_flow_extensions`].
221pub fn should_collect_region_watermark_from_extensions(
222    extensions: &HashMap<String, String>,
223) -> bool {
224    let return_region_seq = extensions
225        .get(FLOW_RETURN_REGION_SEQ)
226        .is_some_and(|value| value.eq_ignore_ascii_case("true"));
227    let has_incremental_after_seqs = extensions.contains_key(FLOW_INCREMENTAL_AFTER_SEQS);
228
229    should_collect_region_watermark(return_region_seq, has_incremental_after_seqs)
230}
231
232fn should_collect_region_watermark(
233    return_region_seq: bool,
234    has_incremental_after_seqs: bool,
235) -> bool {
236    return_region_seq || has_incremental_after_seqs
237}
238
239fn parse_incremental_after_seqs(value: &str) -> Result<HashMap<u64, u64>> {
240    let raw = serde_json::from_str::<HashMap<String, serde_json::Value>>(value).map_err(|e| {
241        invalid_query_context_extension(format!(
242            "Invalid JSON for {}: {} ({})",
243            FLOW_INCREMENTAL_AFTER_SEQS, value, e
244        ))
245    })?;
246
247    raw.into_iter()
248        .map(|(region_id, raw_seq)| {
249            let region_id = region_id.parse::<u64>().map_err(|_| {
250                invalid_query_context_extension(format!(
251                    "Invalid region id in {}: {}",
252                    FLOW_INCREMENTAL_AFTER_SEQS, region_id
253                ))
254            })?;
255
256            let seq = match raw_seq {
257                serde_json::Value::Number(num) => num.as_u64().ok_or_else(|| {
258                    invalid_query_context_extension(format!(
259                        "Invalid sequence value in {} for region {}: {}",
260                        FLOW_INCREMENTAL_AFTER_SEQS, region_id, num
261                    ))
262                })?,
263                serde_json::Value::String(s) => s.parse::<u64>().map_err(|_| {
264                    invalid_query_context_extension(format!(
265                        "Invalid sequence string in {} for region {}: {}",
266                        FLOW_INCREMENTAL_AFTER_SEQS, region_id, s
267                    ))
268                })?,
269                _ => {
270                    return Err(invalid_query_context_extension(format!(
271                        "Invalid sequence value type in {} for region {}",
272                        FLOW_INCREMENTAL_AFTER_SEQS, region_id
273                    )));
274                }
275            };
276
277            Ok((region_id, seq))
278        })
279        .collect()
280}
281
282fn parse_bool(option_name: &str, value: &str) -> Result<bool> {
283    match value {
284        v if v.eq_ignore_ascii_case("true") => Ok(true),
285        v if v.eq_ignore_ascii_case("false") => Ok(false),
286        _ => Err(invalid_query_context_extension(format!(
287            "Invalid value for {}: {}",
288            option_name, value
289        ))),
290    }
291}
292
293/// Parse the scheduled time (in milliseconds since Unix epoch) from extensions.
294///
295/// Returns `Ok(None)` if the extension key is absent, `Ok(Some(millis))` on success,
296/// or `Err` if the value is malformed.
297pub fn parse_scheduled_time_millis(extensions: &HashMap<String, String>) -> Result<Option<i64>> {
298    match extensions.get(FLOW_SCHEDULED_TIME_MILLIS) {
299        Some(val) => val.parse::<i64>().map(Some).map_err(|_| {
300            invalid_query_context_extension(format!(
301                "Invalid value for {}: {}",
302                FLOW_SCHEDULED_TIME_MILLIS, val
303            ))
304        }),
305        None => Ok(None),
306    }
307}
308
309/// Parse the scheduled time from extensions into a [`DateTime<Utc>`].
310///
311/// Returns `Ok(None)` if the extension key is absent, `Ok(Some(datetime))` on success,
312/// or `Err` if the value is malformed.  Millis that produce an out-of-range timestamp are
313/// also rejected.
314pub fn parse_scheduled_time_datetime(
315    extensions: &HashMap<String, String>,
316) -> Result<Option<DateTime<Utc>>> {
317    match parse_scheduled_time_millis(extensions)? {
318        Some(millis) => DateTime::from_timestamp_millis(millis)
319            .map(Some)
320            .ok_or_else(|| {
321                invalid_query_context_extension(format!(
322                    "Out-of-range timestamp for {}: {} ms",
323                    FLOW_SCHEDULED_TIME_MILLIS, millis
324                ))
325            }),
326        None => Ok(None),
327    }
328}
329
330/// Best-effort helper: extract scheduled time from a [`QueryContextRef`] as a [`DateTime<Utc>`].
331///
332/// Errors are silently swallowed, returning `None`; use [`parse_scheduled_time_datetime`]
333/// at call sites that need to reject malformed scheduled time.
334pub fn scheduled_time_from_ctx(query_ctx: &QueryContextRef) -> Option<DateTime<Utc>> {
335    let extensions = query_ctx.extensions();
336    parse_scheduled_time_datetime(&extensions).ok().flatten()
337}
338
339/// Carries the scheduled logical "now" through [`ConfigOptions::extensions`] so
340/// that the distributed plan analyzer can apply it during expression
341/// simplification (preventing wall-clock constant-folding of `now()`).
342#[derive(Debug, Clone)]
343pub(crate) struct ScheduledTimeExtension {
344    pub(crate) scheduled_time: Option<DateTime<Utc>>,
345}
346
347impl ConfigExtension for ScheduledTimeExtension {
348    const PREFIX: &'static str = "flow_scheduled_time";
349}
350
351impl ExtensionOptions for ScheduledTimeExtension {
352    fn as_any(&self) -> &dyn std::any::Any {
353        self
354    }
355
356    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
357        self
358    }
359
360    fn cloned(&self) -> Box<dyn ExtensionOptions> {
361        Box::new(self.clone())
362    }
363
364    fn set(&mut self, key: &str, value: &str) -> datafusion::error::Result<()> {
365        Err(datafusion_common::DataFusionError::NotImplemented(format!(
366            "ScheduledTimeExtension does not support set key: {key} with value: {value}"
367        )))
368    }
369
370    fn entries(&self) -> Vec<ConfigEntry> {
371        vec![]
372    }
373}
374
375fn invalid_query_context_extension(reason: String) -> Error {
376    InvalidQueryContextExtensionSnafu { reason }.build()
377}
378
379#[cfg(test)]
380mod flow_extension_tests {
381    use super::*;
382
383    #[test]
384    fn test_parse_flow_extensions_returns_none_for_non_flow_query() {
385        let exts = HashMap::new();
386        let parsed = FlowQueryExtensions::parse_flow_extensions(&exts).unwrap();
387
388        assert_eq!(parsed, None);
389    }
390
391    #[test]
392    fn test_remote_dyn_filter_pushdown_enabled_from_extensions_defaults_true() {
393        assert!(remote_dyn_filter_pushdown_enabled_from_extensions(&HashMap::new()).unwrap());
394    }
395
396    #[test]
397    fn test_remote_dyn_filter_pushdown_enabled_from_extensions_parses_bool() {
398        let exts = HashMap::from([(
399            QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN.to_string(),
400            "false".to_string(),
401        )]);
402        assert!(!remote_dyn_filter_pushdown_enabled_from_extensions(&exts).unwrap());
403
404        let exts = HashMap::from([(
405            QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN.to_string(),
406            "true".to_string(),
407        )]);
408        assert!(remote_dyn_filter_pushdown_enabled_from_extensions(&exts).unwrap());
409    }
410
411    #[test]
412    fn test_remote_dyn_filter_pushdown_enabled_from_extensions_rejects_invalid_bool() {
413        let exts = HashMap::from([(
414            QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN.to_string(),
415            "invalid".to_string(),
416        )]);
417
418        let err = remote_dyn_filter_pushdown_enabled_from_extensions(&exts).unwrap_err();
419        assert!(format!("{err}").contains(QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN));
420    }
421
422    #[test]
423    fn test_parse_flow_extensions_memtable_only_success() {
424        let exts = HashMap::from([
425            (
426                FLOW_INCREMENTAL_MODE.to_string(),
427                FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY.to_string(),
428            ),
429            (
430                FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
431                r#"{"1":10,"2":20}"#.to_string(),
432            ),
433            (FLOW_RETURN_REGION_SEQ.to_string(), "true".to_string()),
434            (FLOW_SINK_TABLE_ID.to_string(), "1024".to_string()),
435        ]);
436
437        let parsed = FlowQueryExtensions::parse_flow_extensions(&exts)
438            .unwrap()
439            .unwrap();
440        assert_eq!(
441            parsed.incremental_mode,
442            Some(FlowIncrementalMode::MemtableOnly)
443        );
444        assert_eq!(
445            parsed.incremental_after_seqs.unwrap(),
446            HashMap::from([(1, 10), (2, 20)])
447        );
448        assert!(parsed.return_region_seq);
449        assert_eq!(parsed.sink_table_id, Some(1024));
450    }
451
452    #[test]
453    fn test_parse_flow_extensions_mode_requires_after_seqs() {
454        let exts = HashMap::from([(
455            FLOW_INCREMENTAL_MODE.to_string(),
456            FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY.to_string(),
457        )]);
458
459        let err = FlowQueryExtensions::parse_flow_extensions(&exts).unwrap_err();
460        assert!(format!("{err}").contains(FLOW_INCREMENTAL_AFTER_SEQS));
461    }
462
463    #[test]
464    fn test_parse_flow_extensions_invalid_mode() {
465        let exts = HashMap::from([(FLOW_INCREMENTAL_MODE.to_string(), "foo".to_string())]);
466
467        let err = FlowQueryExtensions::parse_flow_extensions(&exts).unwrap_err();
468        assert!(format!("{err}").contains(FLOW_INCREMENTAL_MODE));
469    }
470
471    #[test]
472    fn test_parse_flow_extensions_invalid_after_seqs_json() {
473        let exts = HashMap::from([
474            (
475                FLOW_INCREMENTAL_MODE.to_string(),
476                FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY.to_string(),
477            ),
478            (
479                FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
480                "not-json".to_string(),
481            ),
482        ]);
483
484        let err = FlowQueryExtensions::parse_flow_extensions(&exts).unwrap_err();
485        assert!(format!("{err}").contains(FLOW_INCREMENTAL_AFTER_SEQS));
486    }
487
488    #[test]
489    fn test_parse_flow_extensions_after_seqs_string_values() {
490        let exts = HashMap::from([(
491            FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
492            r#"{"1":"10","2":"20"}"#.to_string(),
493        )]);
494
495        let parsed = FlowQueryExtensions::parse_flow_extensions(&exts)
496            .unwrap()
497            .unwrap();
498        assert_eq!(
499            parsed.incremental_after_seqs.unwrap(),
500            HashMap::from([(1, 10), (2, 20)])
501        );
502    }
503
504    #[test]
505    fn test_parse_flow_extensions_after_seqs_invalid_value_type() {
506        let exts = HashMap::from([(
507            FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
508            r#"{"1":true}"#.to_string(),
509        )]);
510
511        let err = FlowQueryExtensions::parse_flow_extensions(&exts).unwrap_err();
512        assert!(format!("{err}").contains(FLOW_INCREMENTAL_AFTER_SEQS));
513    }
514
515    #[test]
516    fn test_parse_flow_extensions_invalid_sink_table_id() {
517        let exts = HashMap::from([(FLOW_SINK_TABLE_ID.to_string(), "x".to_string())]);
518
519        let err = FlowQueryExtensions::parse_flow_extensions(&exts).unwrap_err();
520        assert!(format!("{err}").contains(FLOW_SINK_TABLE_ID));
521    }
522
523    #[test]
524    fn test_validate_for_scan_missing_source_region() {
525        let source_region_id = RegionId::new(100, 2);
526        let existing_region_id = RegionId::new(100, 1);
527        let exts = HashMap::from([
528            (
529                FLOW_INCREMENTAL_MODE.to_string(),
530                FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY.to_string(),
531            ),
532            (
533                FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
534                format!(r#"{{"{}":10}}"#, existing_region_id.as_u64()),
535            ),
536        ]);
537
538        let parsed = FlowQueryExtensions::parse_flow_extensions(&exts)
539            .unwrap()
540            .unwrap();
541        let err = parsed.validate_for_scan(source_region_id).unwrap_err();
542        assert!(format!("{err}").contains("Missing region"));
543    }
544
545    #[test]
546    fn test_validate_for_scan_sink_table_excluded() {
547        let source_region_id = RegionId::new(1024, 1);
548        let exts = HashMap::from([
549            (
550                FLOW_INCREMENTAL_MODE.to_string(),
551                FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY.to_string(),
552            ),
553            (
554                FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
555                format!(r#"{{"{}":10}}"#, source_region_id.as_u64()),
556            ),
557            (FLOW_SINK_TABLE_ID.to_string(), "1024".to_string()),
558        ]);
559
560        let parsed = FlowQueryExtensions::parse_flow_extensions(&exts)
561            .unwrap()
562            .unwrap();
563        let apply_incremental = parsed.validate_for_scan(source_region_id).unwrap();
564        assert!(!apply_incremental);
565    }
566
567    #[test]
568    fn test_should_collect_region_watermark_defaults_false() {
569        let parsed = FlowQueryExtensions::default();
570        assert!(!parsed.should_collect_region_watermark());
571    }
572
573    #[test]
574    fn test_should_collect_region_watermark_true_for_return_region_seq() {
575        let parsed = FlowQueryExtensions {
576            return_region_seq: true,
577            ..Default::default()
578        };
579        assert!(parsed.should_collect_region_watermark());
580    }
581
582    #[test]
583    fn test_should_collect_region_watermark_true_for_incremental_query() {
584        let parsed = FlowQueryExtensions {
585            incremental_after_seqs: Some(HashMap::from([(1, 10)])),
586            ..Default::default()
587        };
588        assert!(parsed.should_collect_region_watermark());
589    }
590
591    #[test]
592    fn test_should_collect_region_watermark_from_extensions() {
593        let exts = HashMap::from([(FLOW_RETURN_REGION_SEQ.to_string(), "true".to_string())]);
594        assert!(should_collect_region_watermark_from_extensions(&exts));
595
596        let exts = HashMap::from([(
597            FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
598            r#"{"1":10}"#.to_string(),
599        )]);
600        assert!(should_collect_region_watermark_from_extensions(&exts));
601
602        let exts = HashMap::from([(FLOW_RETURN_REGION_SEQ.to_string(), "false".to_string())]);
603        assert!(!should_collect_region_watermark_from_extensions(&exts));
604        assert!(!should_collect_region_watermark_from_extensions(
605            &HashMap::new()
606        ));
607    }
608
609    #[test]
610    fn test_parse_flow_extensions_return_region_seq_only_returns_some() {
611        let exts = HashMap::from([(FLOW_RETURN_REGION_SEQ.to_string(), "true".to_string())]);
612
613        let parsed = FlowQueryExtensions::parse_flow_extensions(&exts)
614            .unwrap()
615            .unwrap();
616
617        assert!(parsed.return_region_seq);
618    }
619
620    #[test]
621    fn test_parse_flow_extensions_sink_table_only_returns_some() {
622        let exts = HashMap::from([(FLOW_SINK_TABLE_ID.to_string(), "1024".to_string())]);
623
624        let parsed = FlowQueryExtensions::parse_flow_extensions(&exts)
625            .unwrap()
626            .unwrap();
627
628        assert_eq!(parsed.sink_table_id, Some(1024));
629    }
630
631    #[test]
632    fn test_parse_flow_extensions_incremental_after_seqs_only_returns_some() {
633        let exts = HashMap::from([(
634            FLOW_INCREMENTAL_AFTER_SEQS.to_string(),
635            r#"{"1":10}"#.to_string(),
636        )]);
637
638        let parsed = FlowQueryExtensions::parse_flow_extensions(&exts)
639            .unwrap()
640            .unwrap();
641
642        assert_eq!(
643            parsed.incremental_after_seqs,
644            Some(HashMap::from([(1, 10)]))
645        );
646    }
647
648    // --- scheduled time helper tests ---
649
650    #[test]
651    fn test_parse_scheduled_time_millis_absent() {
652        let exts = HashMap::new();
653        assert_eq!(parse_scheduled_time_millis(&exts).unwrap(), None);
654    }
655
656    #[test]
657    fn test_parse_scheduled_time_millis_valid() {
658        let exts = HashMap::from([(
659            FLOW_SCHEDULED_TIME_MILLIS.to_string(),
660            "1700000000000".to_string(),
661        )]);
662        assert_eq!(
663            parse_scheduled_time_millis(&exts).unwrap(),
664            Some(1700000000000)
665        );
666    }
667
668    #[test]
669    fn test_parse_scheduled_time_millis_malformed() {
670        let exts = HashMap::from([(
671            FLOW_SCHEDULED_TIME_MILLIS.to_string(),
672            "not-a-number".to_string(),
673        )]);
674        let err = parse_scheduled_time_millis(&exts).unwrap_err();
675        assert!(format!("{err}").contains(FLOW_SCHEDULED_TIME_MILLIS));
676    }
677
678    #[test]
679    fn test_parse_scheduled_time_datetime_valid() {
680        let exts = HashMap::from([(
681            FLOW_SCHEDULED_TIME_MILLIS.to_string(),
682            "1700000000000".to_string(),
683        )]);
684        let dt = parse_scheduled_time_datetime(&exts).unwrap().unwrap();
685        assert_eq!(dt.timestamp_millis(), 1700000000000);
686    }
687
688    #[test]
689    fn test_parse_scheduled_time_datetime_negative_millis() {
690        let exts = HashMap::from([(FLOW_SCHEDULED_TIME_MILLIS.to_string(), "-1".to_string())]);
691        let dt = parse_scheduled_time_datetime(&exts).unwrap().unwrap();
692        assert_eq!(dt.timestamp_millis(), -1);
693    }
694
695    #[test]
696    fn test_parse_scheduled_time_datetime_out_of_range() {
697        // i64::MAX millis is well beyond representable chrono::DateTime range
698        let exts = HashMap::from([(FLOW_SCHEDULED_TIME_MILLIS.to_string(), i64::MAX.to_string())]);
699        let err = parse_scheduled_time_datetime(&exts).unwrap_err();
700        assert!(format!("{err}").contains("Out-of-range"));
701    }
702}
query/options.rs

query/
options.rs
