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

query/
options.rs
